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 = "HuggingFaceTB/SmolLM-1.7B"
+assistant_id = "HuggingFaceTB/SmolLM-135M"
 
-model = AutoModelForCausalLM.from_pretrained(model_id, attn_implementation="eager").to(dtype=torch.bfloat16, device="cuda")
-assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id, attn_implementation="eager").to(dtype=torch.bfloat16, device="cuda")
+model = AutoModelForCausalLM.from_pretrained(model_id).to(dtype=torch.bfloat16, device="cuda")
+assistant_model = AutoModelForCausalLM.from_pretrained(assistant_id).to(dtype=torch.bfloat16, device="cuda")
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 @spaces.GPU