chore: import upstream snapshot with attribution

gpu/stats.py

@@ -0,0 +1,57 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
+#
+# This source code is licensed under the BSD license found in the
+# LICENSE file in the root directory of this source tree.
+
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass
+class PhaseStats:
+    name: str
+    tokens: int
+    time: float
+
+    def show(self) -> str:
+        tps = self.tokens / self.time
+        return (
+            f"[{self.name}] "
+            f"generated tokens: {self.tokens}"
+            f" - total time: {self.time:.3f}s"
+            f" - {tps:.1f} tokens per second"
+        )
+
+
+class Stats:
+    """
+    Generation stats, split by phases.
+    """
+
+    def __init__(self):
+        self.phases = []
+        self.current = None
+
+    def end_phase(self, tokens: int, now: Optional[float] = None):
+        """Terminate the current phase."""
+        if self.current is None:
+            return
+        if now is None:
+            now = time.time()
+        cname, ctokens, ctime = self.current
+        stats = PhaseStats(
+            name=cname,
+            tokens=tokens - ctokens,
+            time=now - ctime,
+        )
+        self.phases.append(stats)
+
+    def phase(self, name: str, tokens: int = 0):
+        """
+        Start a new phase, and terminate the current one,
+        if one is ongoing.
+        """
+        now = time.time()
+        self.end_phase(tokens, now)
+        self.current = (name, tokens, now)