Update app.py

app.py

@@ -17,27 +17,25 @@ model = AutoModelForCausalLM.from_pretrained(
 # --- Chat function ---
 def chat(message, history, max_new_tokens=128, temperature=0.7):
     try:
-        # Convert history into Hugging Face messages format
+        # Convert conversation history into messages
         messages = []
         for user_msg, bot_msg in history:
             messages.append({"role": "user", "content": user_msg})
             messages.append({"role": "assistant", "content": bot_msg})
         messages.append({"role": "user", "content": message})
 
-        # Prepare inputs with chat template → return dictionary
-        inputs = tokenizer.apply_chat_template(
+        # Apply chat template -> returns tensor of input_ids
+        input_ids = tokenizer.apply_chat_template(
             messages,
             add_generation_prompt=True,
             tokenize=True,
-            return_tensors="pt",
-            return_dict=True
-        )
+            return_tensors="pt"
+        ).to(model.device)
 
-        # Move all tensors in the input dict to the model device
-        for k in inputs:
-            inputs[k] = inputs[k].to(model.device)
+        # Wrap as dict so generate(**inputs) works
+        inputs = {"input_ids": input_ids}
 
-        # Generate model output
+        # Generate output
         outputs = model.generate(
             **inputs,
             max_new_tokens=max_new_tokens,
@@ -46,19 +44,18 @@ def chat(message, history, max_new_tokens=128, temperature=0.7):
             pad_token_id=tokenizer.eos_token_id
         )
 
-        # Decode ONLY the newly generated tokens (past the input length)
+        # Decode new tokens only
         response = tokenizer.decode(
-            outputs[0][inputs["input_ids"].shape[-1]:],
+            outputs[0][input_ids.shape[-1]:],
             skip_special_tokens=True
         ).strip()
 
-        # Append to history
         history.append((message, response))
         return history, history, ""
 
     except Exception as e:
         import traceback
-        traceback.print_exc()  # will show in HF Space Logs
+        traceback.print_exc()
         return history + [(message, f"⚠️ Error: {str(e)}")], history, ""