Update app.py

app.py

@@ -28,7 +28,6 @@ print("Model ready ✓")
 
 # ── Inference ─────────────────────────────────────────────────────────────────
 def respond(message: str, history: list, system_prompt: str, max_new_tokens: int):
-    """Build the chat prompt and return a response."""
 
     messages = []
 
@@ -37,7 +36,6 @@ def respond(message: str, history: list, system_prompt: str, max_new_tokens: int
     if system_prompt.strip() and not history:
         effective_message = f"{system_prompt.strip()}\n\n{message}"
 
-    # history is a list of dicts with 'role' and 'content' keys (messages format)
     for turn in history:
         messages.append({"role": turn["role"], "content": turn["content"]})
 
@@ -50,6 +48,10 @@ def respond(message: str, history: list, system_prompt: str, max_new_tokens: int
         return_tensors="pt",
     )
 
+    # Gemma 4 returns a BatchEncoding dict instead of a raw tensor
+    if not isinstance(input_ids, torch.Tensor):
+        input_ids = input_ids["input_ids"]
+
     # Generate
     with torch.no_grad():
         output_ids = model.generate(