Update app.py

app.py

@@ -8,11 +8,11 @@ import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 
-model_id = "facebook/opt-6.7b"
-assistant_id = "facebook/opt-125m"
+model_id = "Qwen/Qwen2.5-32B-Instruct"
+assistant_id = "Qwen/Qwen2.5-0.5B-Instruct"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, device_map="auto")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
+model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=True, device_map="auto")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.float16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
@@ -49,9 +49,9 @@ def run_generation(user_text, use_assistant, temperature, max_new_tokens):
     model_output = ""
     for new_text in streamer:
         model_output += new_text
-        time_so_far = round(time.time() - start, 3)
+        time_so_far = time.time() - start
         tokens_so_far = tokenizer(model_output, return_tensors="pt").input_ids.shape[1]
-        yield [model_output, tokens_so_far/time_so_far]
+        yield [model_output, round(tokens_so_far/time_so_far, 2)]
 
 
 def reset_textbox():
@@ -61,8 +61,8 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"- Model: {model_id} (INT8, ~7GB)\n"
-        f"- Assistant Model: {assistant_id} (FP16, ~0.3GB)\n"
+        f"- Model: {model_id} (4-bit quant, ~16GB)\n"
+        f"- Assistant Model: {assistant_id} (FP16, ~1GB)\n"
         "- Recipe for speedup: a) >10x model size difference in parameters; b) assistant trained similarly; c) CPU is not a bottleneck"
     )
 
@@ -84,7 +84,7 @@ with gr.Blocks() as demo:
             temperature = gr.Slider(
                 minimum=0.0, maximum=2.0, value=0.6, step=0.05, interactive=True, label="Temperature (0.0 = Greedy)",
             )
-            gr.Markdown("### Tokens per secon")
+            gr.Markdown("### Tokens per second")
             tokens_per_second = gr.Textbox(lines=1, interactive=False, show_label=False)
 
     generate_inputs = [user_text, use_assistant, temperature, max_new_tokens]