| 48 | |
| 49 | |
| 50 | def stream_output(output_stream): |
| 51 | print(f"ASSISTANT: ", end="", flush=True) |
| 52 | pre = 0 |
| 53 | for outputs in output_stream: |
| 54 | output_text = outputs["text"] |
| 55 | output_text = output_text.strip().split(" ") |
| 56 | now = len(output_text) - 1 |
| 57 | if now > pre: |
| 58 | print(" ".join(output_text[pre:now]), end=" ", flush=True) |
| 59 | pre = now |
| 60 | print(" ".join(output_text[pre:]), flush=True) |
| 61 | if "timing" in outputs and outputs["timing"] is not None: |
| 62 | timing = outputs["timing"] |
| 63 | context_tokens = timing["context_tokens"] |
| 64 | context_time = timing["context_time"] |
| 65 | total_tokens = timing["total_tokens"] |
| 66 | generation_time_list = timing["generation_time_list"] |
| 67 | generation_tokens = len(generation_time_list) |
| 68 | average_speed = (context_time + np.sum(generation_time_list)) / ( |
| 69 | context_tokens + generation_tokens |
| 70 | ) |
| 71 | print("=" * 50) |
| 72 | print("Speed of Inference") |
| 73 | # print("-" * 50) |
| 74 | # print(f"Context Stage : {context_time/context_tokens * 1000:.2f} ms/token") |
| 75 | print("-" * 50) |
| 76 | print( |
| 77 | f"Generation Stage : {np.average(generation_time_list) * 1000:.2f} ms/token" |
| 78 | ) |
| 79 | # print(f"Average Speed : {average_speed * 1000:.2f} ms/token") |
| 80 | print("=" * 50) |
| 81 | # print("token num:", total_tokens) |
| 82 | # print("Model total Time = ", (context_time + np.sum(generation_time_list))*1000, "ms" ) |
| 83 | return " ".join(output_text) |
| 84 | |
| 85 | |
| 86 | if __name__ == "__main__": |