Update handler.py

handler.py

@@ -157,11 +157,10 @@ Note - each answer must be of length <= 400."""
                     **gen_config
                 )
             
-            # Decode response
-            full_response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-            
-            # Extract only the assistant's response
-            assistant_response = full_response.split("<|start_header_id|>assistant<|end_header_id|>")[-1].strip()
+            # Decode response - only decode the generated tokens, not the input
+            input_length = model_inputs['input_ids'].shape[1]
+            generated_tokens = outputs[0][input_length:]
+            assistant_response = self.tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
             
             # Extract the final answer
             extracted_answer = self.extract_answer(assistant_response)