Update app.py

app.py

@@ -21,10 +21,22 @@ MODEL_ID = "google/gemma-4-31B-it-assistant"
 print("Loading tokenizer...")
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
 
+# Manually set Gemma chat template since tokenizer doesn't include one
+tokenizer.chat_template = (
+    "{% for message in messages %}"
+    "{% if message['role'] == 'user' %}"
+    "user\n{{ message['content'] }}\n"
+    "{% elif message['role'] == 'assistant' %}"
+    "model\n{{ message['content'] }}\n"
+    "{% endif %}"
+    "{% endfor %}"
+    "{% if add_generation_prompt %}model\n{% endif %}"
+)
+
 print("Loading model...")
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.bfloat16,
+    dtype=torch.bfloat16,
     device_map="auto",
 )
 model.eval()
@@ -73,11 +85,6 @@ demo = gr.ChatInterface(
     fn=chat,
     title="Gemma 4 Assistant",
     description="google/gemma-4-31B-it-assistant — streaming enabled",
-    examples=[
-        "Explain quantum computing in simple terms",
-        "Write a Python function to reverse a string",
-        "What is photosynthesis?",
-    ],
 )
 
 demo.launch()