Add chat template

app.py

@@ -12,7 +12,10 @@ MODEL_ID = "le-llm/gemma-3-12b-it-reasoning"
 # Load model & tokenizer
 device = "cuda" if torch.cuda.is_available() else "cpu"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype=torch.float16 if device=="cuda" else torch.float32).to(device)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32
+).to(device)
 
 SYSTEM_PROMPT = "You are a friendly Chatbot."
 
@@ -24,14 +27,17 @@ def respond(
     temperature,
     top_p,
 ):
+    # Build conversation in chat template format
+    messages = [{"role": "system", "content": system_message}] + history + [
+        {"role": "user", "content": message}
+    ]
 
-    conversation = system_message + "\n"
-    for turn in history:
-        role = "User" if turn["role"] == "user" else "Assistant"
-        conversation += f"{role}: {turn['content']}\n"
-    conversation += f"User: {message}\nAssistant:"
-
-    inputs = tokenizer(conversation, return_tensors="pt").to(device)
+    input_text = tokenizer.apply_chat_template(
+        messages,
+        tokenize=False,
+        add_generation_prompt=True  # ensures model knows it's assistant's turn
+    )
+    inputs = tokenizer(input_text, return_tensors="pt").to(device)
 
     output_ids = model.generate(
         **inputs,
@@ -39,10 +45,15 @@ def respond(
         temperature=temperature,
         top_p=top_p,
         do_sample=True,
+        eos_token_id=tokenizer.eos_token_id,  # stop at EOS
     )
 
-    response = tokenizer.decode(output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
-    yield response
+    # Only return the newly generated assistant message
+    response = tokenizer.decode(
+        output_ids[0][inputs["input_ids"].shape[1]:],
+        skip_special_tokens=True
+    )
+    return response
 
 chatbot = gr.ChatInterface(
     respond,
@@ -61,4 +72,4 @@ chatbot = gr.ChatInterface(
     ],
 )
 
-chatbot.launch()
+chatbot.launch()