Update handler.py

handler.py

@@ -4,8 +4,11 @@ from typing import Dict, List, Any
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-MAX_TOKENS=8192
-GPU_LAYERS=99 if torch.cuda.is_available() else 0
+PROMPT_FORMAT= """
+<|user|>
+{inputs} <|end|>
+<|assistant|>
+"""
 
 class EndpointHandler():
     def __init__(self, data):
@@ -17,19 +20,20 @@ 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_length = data.pop("max_length", 1024)
+        max_new_tokens = data.pop("max_length", 1024)
         try:
-            max_length = int(max_length)
+            max_new_tokens = int(max_new_tokens)
         except Exception as e:
             return json.dumps({
                 "status": "error",
                 "reason": "max_length was passed as something that was absolutely not a plain old int"
             })
         
-        res = self.model(f"""
-<|user|>
-{inputs} <|end|>
-<|assistant|>
-""", max_new_tokens=max_new_tokens, do_sample=False)
+        res = PROMPT_FORMAT.format(do_sample=False)
+        retrurn model(
+            res,
+            max_new_tokens=max_new_tokens,
+            do_sample=False
+        )
 
         return res