Update gen.py

gen.py

@@ -6,6 +6,7 @@ import sys
 # Cargar el pipeline de generación de texto
 pipe = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-alpha", torch_dtype=torch.bfloat16, device_map="auto")
 
+
 # Definir el prompt para generar un JSON con eventos anidados
 prompt = (
     "Genera un JSON que describa una serie de eventos consecutivos en un formato similar al siguiente:\n\n"
@@ -148,7 +149,16 @@ prompt = (
 
 def generate(event):
     # Generar el texto usando el modelo
-    outputs = pipe(prompt+" "+event, max_new_tokens=512, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
+    # We use the tokenizer's chat template to format each message - see https://huggingface.co/docs/transformers/main/en/chat_templating
+    messages = [
+        {
+            "role": "system",
+            "content": prompt,
+        },
+        {"role": "user", "content": event},
+    ]
+    prompt_tuned = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    outputs = pipe(prompt_tuned, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
 
     # Imprimir la salida generada
     return outputs