actually do rolling window

PrimeIntellect-ai · Sep 26, 2024 · 867350d · 867350d
1 parent 0dc5638
commit 867350d
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/zeroband/utils/__init__.py b/src/zeroband/utils/__init__.py
@@ -82,6 +82,9 @@ def __init__(self, window_size: int):
     def count_tokens(self, tokens: int):
         self.tokens.append(tokens)
         self.times.append(time.perf_counter())
+        if len(self.tokens) > self.window_size:
+            self.tokens.pop(0)
+            self.times.pop(0)
 
     def get_tokens_per_second(self) -> float | None:
         if len(self.tokens) < 2: