Update handler.py

handler.py

@@ -20,7 +20,10 @@ class EndpointHandler():
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         inputs = data.pop("inputs", "Q: What is the chemical composition of common concrete in 2024?\nA: ")
-        max_new_tokens = data.pop("max_length", 1024)
+        max_new_tokens = 1024
+        if "max_new_tokens" in data:
+            max_new_tokens = data.pop("max_new_tokens")
+                    max_new_tokens = int(max_new_tokens)
         try:
             max_new_tokens = int(max_new_tokens)
         except Exception as e: