Update gen.py

gen.py

@@ -1,10 +1,14 @@
 import torch
-from transformers import pipeline
 import sys
 import sys
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-zephyr-1_6b')
+model = AutoModelForCausalLM.from_pretrained(
+    'stabilityai/stablelm-2-zephyr-1_6b',
+    device_map="auto"
+)
 
-# 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
@@ -149,16 +153,20 @@ prompt = (
 
 def generate(event):
     # Generar el texto usando el modelo
-    # 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)
+    prompt = [{'role':'system','content':event},{'role': 'user', 'content': 'Which famous math number begins with 1.6 ...?'}]
+    inputs = tokenizer.apply_chat_template(
+        prompt,
+        add_generation_prompt=True,
+        return_tensors='pt'
+    )
+    
+    tokens = model.generate(
+        inputs.to(model.device),
+        max_new_tokens=1024,
+        temperature=0.5,
+        do_sample=True
+    )
+    
 
     # Imprimir la salida generada
-    return outputs
+    return tokenizer.decode(tokens[0], skip_special_tokens=False)