Update app.py

app.py

@@ -6,10 +6,10 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStream
 import time
 
 model_id = "facebook/opt-6.7b"
-assistant_id = "facebook/opt-350m"
+assistant_id = "facebook/opt-125m"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(device=model.device, dtype=torch.bfloat16)
+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.float16)
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 def run_generation(user_text, use_assistant, temperature, max_new_tokens):
@@ -55,8 +55,8 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"- Model: {model_id} (BF16, ~14GB)\n"
-        f"- Assistant Model: {assistant_id} (BF16, ~0.3GB)\n"
+        f"- Model: {model_id} (FP16, ~14GB)\n"
+        f"- Assistant Model: {assistant_id} (FP16, ~0.3GB)\n"
         "- Running on a T4 GPU\n"
         "- Best results are obtained when the model size difference in parameters is >10x"
     )