Normalize TinyLlama chat template inputs

app.py

@@ -576,8 +576,15 @@ def _generate_tinyllama(prompt, max_new_tokens=64, temperature=0.7):
     temperature = max(0.0, min(2.0, float(temperature)))
     messages = [{"role": "user", "content": prompt}]
     if hasattr(tokenizer, "apply_chat_template"):
-        inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
-        attention_mask = torch.ones_like(inputs)
+        templated = tokenizer.apply_chat_template(
+            messages, add_generation_prompt=True, return_tensors="pt"
+        )
+        if isinstance(templated, dict) or hasattr(templated, "input_ids"):
+            inputs = templated["input_ids"]
+            attention_mask = templated.get("attention_mask", torch.ones_like(inputs))
+        else:
+            inputs = templated
+            attention_mask = torch.ones_like(inputs)
     else:
         encoded = tokenizer(prompt, return_tensors="pt")
         inputs = encoded["input_ids"]
@@ -585,7 +592,7 @@ def _generate_tinyllama(prompt, max_new_tokens=64, temperature=0.7):
     t0 = time.time()
     with torch.inference_mode():
         output = model.generate(
-            inputs,
+            input_ids=inputs,
             attention_mask=attention_mask,
             max_new_tokens=max_new_tokens,
             do_sample=temperature > 0,