Spaces:
Sleeping
Sleeping
| """ | |
| Token generation speed measurement. | |
| TokenSpeedometer tracks wall-clock time across a generation run and | |
| reports tokens-per-second (tok/s). Useful for benchmarking different | |
| model sizes, devices, or sampling strategies. Call `start` before | |
| the loop, `tick` after each token, and read `tps` for the running | |
| average. | |
| """ | |
| import time | |
| class TokenSpeedometer: | |
| def __init__(self): | |
| self.start_time = None | |
| self.first_token_time = None | |
| self.end_time = None | |
| self.token_count = 0 | |
| def start(self): | |
| """Call this right before the generation loop starts.""" | |
| self.start_time = time.perf_counter() | |
| return self | |
| def update(self): | |
| """Call this inside the loop every time a new token is yielded.""" | |
| if self.token_count == 0: | |
| self.first_token_time = time.perf_counter() | |
| self.token_count += 1 | |
| def end(self): | |
| """Call this after the loop finishes.""" | |
| self.end_time = time.perf_counter() | |
| self.print_stats() | |
| def print_stats(self): | |
| if self.token_count == 0: | |
| print("No tokens generated.") | |
| return | |
| total_duration = self.end_time - self.start_time | |
| # 1. Overall Speed (including prompt processing) | |
| overall_tps = self.token_count / total_duration | |
| print(f"\n--- Performance Metrics ---") | |
| print(f"Total Tokens: {self.token_count}") | |
| print(f"Total Time: {total_duration:.2f} s") | |
| print(f"Overall Speed: {overall_tps:.2f} tokens/sec") | |
| # 2. Decoding Speed (excluding the wait for the first token) | |
| # This measures purely how fast the model generates text once it starts | |
| if self.token_count > 1 and self.first_token_time: | |
| decoding_duration = self.end_time - self.first_token_time | |
| decoding_tokens = self.token_count - 1 | |
| decoding_tps = decoding_tokens / decoding_duration | |
| ttft = (self.first_token_time - self.start_time) * 1000 # Time To First Token | |
| print(f"Decoding Speed: {decoding_tps:.2f} tokens/sec (excluding pre-fill)") | |
| print(f"Time to 1st Tok: {ttft:.2f} ms") | |
| print("---------------------------") |