gstaff/token-per-second-simulator

1import gradio as gr2import time3from transformers import AutoTokenizer4 5tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")6 7starter_text = """# Abstract8 Within thirty years, we will have the technological means to create superhuman intelligence. Shortly after,9 the human era will be ended.10 Is such progress avoidable? If not to be avoided, can events be guided so that we may survive? These questions11 are investigated. Some possible answers (and some further dangers) are presented.12"""13 14 15def calculate_wait_seconds(tokens_per_second):16    return 1 / tokens_per_second17 18 19def get_tokens(prompt):20    tokens = tokenizer.tokenize(prompt)21    return [x.replace('▁', ' ').replace('<0x0A>', '\n') for x in tokens]22 23 24def echo(message, history, prompt, tokens_per_second, time_to_first_token, stream):25    wait_seconds = calculate_wait_seconds(tokens_per_second)26 27    response = f"{prompt}"28    tokens = get_tokens(response)29 30    if time_to_first_token:31        time.sleep(time_to_first_token / 1000)32    partial_message = ""33    for new_token in tokens:34        time.sleep(wait_seconds)35        if '<' in new_token:36            # Gradio chat chokes on HTML-like elements37            continue38        partial_message += str(new_token)39        if stream:40            yield partial_message41 42    if not stream:43        yield partial_message44 45 46with gr.Blocks(title='Tokens per Second Simulator') as demo:47    gr.Markdown('# ⏱️ Tokens per Second Simulator')48    gr.Markdown('Compare the feel of different response speeds for a chat bot')49    gr.Markdown('Reading speeds vary but in English 5-10 tokens per second is considered normal reading speed')50    gr.Markdown(51        'References for further research:\n'52        '- https://www.perplexity.ai/search/How-many-tokens-1d7VyXCDQuWf3pJnK4.0iw?s=c\n'53        '- https://www.databricks.com/blog/llm-inference-performance-engineering-best-practices\n'54        '- https://news.ycombinator.com/item?id=35978864\n'55        '- https://www.reddit.com/r/LocalLLaMA/comments/162pgx9/what_do_yall_consider_acceptable_tokens_per/')56 57    prompt = gr.Textbox(starter_text, label="Prompt to Echo")58    tps_slider = gr.Slider(1, 50, render=True, value=8, label='Tokens per second (TPS)')59    ttft_slider = gr.Slider(0, 5000, render=True, value=0,60                            label='Time to first token (TTFT) in milliseconds')61    stream_checkbox = gr.Checkbox(label='Stream Response', value=True)62 63    gr.ChatInterface(echo, additional_inputs=[prompt, tps_slider, ttft_slider, stream_checkbox],64                     description='Submit any text to echo the prompt above at the selected speed.')65 66demo.queue().launch()67