Update handler.py

handler.py

@@ -3,19 +3,11 @@ from typing import Dict, Any, List
 import torch
 import json
 
-def model_inference(model, tokenizer, text, max_new_tokens, temperature, top_p=0.99, top_k=20, repetition_penalty=1.2):
-  inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
-  inputs_length = len(inputs["input_ids"][0])
-  with torch.inference_mode():
-    outputs = model.generate(**inputs,
-                             max_new_tokens=max_new_tokens,
-                             temperature=temperature,
-                             top_p=top_p,
-                             top_k=top_k,
-                             repetition_penalty=repetition_penalty)
-  return tokenizer.decode(outputs[0][inputs_length:], skip_special_tokens=True)
-
-
+def model_inference(model, tokenizer, text, **kwargs):
+    inputs = tokenizer(text, return_tensors="pt").to("cuda:0")
+    with torch.inference_mode():
+        outputs = model.generate(**inputs, **kwargs)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 class EndpointHandler():
     def __init__(self, path=""):
@@ -27,24 +19,17 @@ class EndpointHandler():
         )
         self.my_tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
 
-    def __call__(self, data):
-        print(data)
+    def __call__(self, data: Dict[str, Any]) -> Dict[str, str]:
         input_data = data.pop("inputs", data)
+        parameters = data.pop("parameters", {})
         template = """
-        [INST] <<SYS>>
+        <s>[INST] <<SYS>>
         A continuación se presentan parámetros fisiológicos y ambientales capturados por los sensores de un smartwatch utilizado por un atleta. Basándote en estos datos y en tu conocimiento de las rutinas de entrenamiento para atletas, crea una rutina diaria de entrenamiento personalizada para el atleta.
         <</SYS>>
         {data} [/INST]
         """.strip()
-
         data_text = json.dumps(input_data, indent=4)
         filled_template = template.replace("{data}", data_text)
-        print("input:")
-        print(filled_template)
-        max_new_tokens = 1600
-        temperature = 0.5
+        result = model_inference(self.model, self.my_tokenizer, filled_template, **parameters)
         
-        result = model_inference(self.model, self.my_tokenizer, filled_template, max_new_tokens, temperature)
-        print("output:")
-        print(result)
         return {'generated_text': [result]}