Update app.py

app.py

@@ -6,11 +6,11 @@ import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 import time
 
-model_id = "google/gemma-2-9b-it"
-assistant_id = "google/gemma-2-2b-it"
+model_id = "facebook/opt-1.3b"
+assistant_id = "facebook/opt-125m"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto", load_in_8bit=True, attn_implementation="eager")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, torch_dtype=torch.bfloat16, device_map="auto", attn_implementation="eager")
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, torch_dtype=torch.bfloat16, device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 @spaces.GPU
@@ -56,9 +56,9 @@ def reset_textbox():
 with gr.Blocks() as demo:
     gr.Markdown(
         "# 🤗 Assisted Generation Demo\n"
-        f"- Model: {model_id} (quantized to 8 bits with BNB, ~9GB)\n"
-        f"- Assistant Model: {assistant_id} (BF16, ~4GB)\n"
-        "- Running on a T4 GPU"
+        f"- Model: {model_id} (BF16, ~3GB)\n"
+        f"- Assistant Model: {assistant_id} (BF16, ~0.3GB)\n"
+        "- Running on a T4 GPU\n"
         "- Best results are obtained when the model size difference in parameters is >10x"
     )