Update handler.py

handler.py

@@ -25,24 +25,34 @@ class EndpointHandler:
         prompt = data.get("inputs", "")
         if not prompt:
             return {"error": "Missing 'inputs' field."}
-        
+
         defaults = {
             "max_new_tokens": 100,
             "do_sample": True,
             "temperature": 0.7,
-            "top_p": 0.9
+            "top_p": 0.9,
+            "eos_token_id": self.tokenizer.eos_token_id  # ✅ Stop at <|eot_id|>
         }
-    
+
         generation_args = {**defaults, **data.get("parameters", {})}
-    
-        outputs = self.generator(prompt, **generation_args)
-    
-        return {
-            "choices": [{
-                "message": {
-                    "role": "assistant",
-                    "content": outputs[0]["generated_text"].strip()
-                },
-                "finish_reason": "stop"
-            }]
-        }
+
+        try:
+            outputs = self.generator(prompt, **generation_args)
+            output_text = outputs[0]["generated_text"].strip()
+
+            finish_reason = "stop"
+            if len(self.tokenizer.encode(output_text)) >= generation_args["max_new_tokens"]:
+                finish_reason = "length"
+
+            return {
+                "choices": [{
+                    "message": {
+                        "role": "assistant",
+                        "content": output_text
+                    },
+                    "finish_reason": finish_reason
+                }]
+            }
+
+        except Exception as e:
+            return {"error": str(e)}