Update handler.py

handler.py

@@ -4,8 +4,8 @@ import json
 
 class EndpointHandler:
     def __init__(self, model_dir):
-        self.tokenizer = AutoTokenizer.from_pretrained("facebook/m2m100_418M")
-        self.model = AutoModelForSeq2SeqLM.from_pretrained("facebook/m2m100_418M")
+        self.tokenizer = AutoTokenizer.from_pretrained("jla25/squareV3")
+        self.model = AutoModelForSeq2SeqLM.from_pretrained("jla25/squareV3")
         self.model.eval()
 
     def preprocess(self, data):
@@ -15,37 +15,23 @@ class EndpointHandler:
         # Prompt personalizado para guiar al modelo
         input_text = (
             f"""
-Genera un JSON válido con estas especificaciones:
-- Cada objeto tiene una clave 'id' y un valor 'value'.
-- Opciones para 'id': firstName, lastName, jobTitle, address, email, phone, notes, roleFunction.
-- Si 'id' es address, email o phone, debe incluir subclaves: MOBILE, WORK, PERSONAL, MAIN, OTHER.
-- 'roleFunction' debe ser una de estas: BUYER, SELLER, SUPPLIER, PARTNER, COLLABORATOR, PROVIDER, CUSTOMER.
-Ejemplo:
-Entrada: "Contacté a Juan Pérez, Gerente de Finanzas."
-Salida esperada:
-{{
-    "values": [
-        {{"id": "firstName", "value": "Juan"}},
-        {{"id": "lastName", "value": "Pérez"}},
-        {{"id": "jobTitle", "value": "Gerente de Finanzas"}}
-    ]
-}}
-        Procesa este texto: "{data['inputs']}"
+          ### Procesa el siguiente texto y genera un JSON válido:
+          "{data['inputs']}"
             """)
         # Imprimir el texto generado para el prompt
         print(f"Prompt generado para el modelo: {input_text}")
-        tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=1000)
+        tokens = self.tokenizer(input_text, return_tensors="pt", truncation=True, padding="max_length", max_length=1024)
         return tokens
 
     def inference(self, tokens):
         generate_kwargs = {
-            "max_length": 1500,
-            "num_beams": 7,
+            "max_length": 1024,
+            "num_beams": 5,
             "do_sample": False,
-            "temperature": 0.1,
-            "top_k": 10,
+            "temperature": 0.3,
+            "top_k": 50,
             "top_p": 0.7,
-            "repetition_penalty": 2.8
+            "repetition_penalty": 2.5
         }
         with torch.no_grad():
             outputs = self.model.generate(**tokens, **generate_kwargs)