Update gen.py

gen.py

@@ -3,10 +3,10 @@ import sys
 import sys
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-tokenizer = AutoTokenizer.from_pretrained('google/gemma-2-2b-it')
+tokenizer = AutoTokenizer.from_pretrained('stabilityai/stablelm-2-zephyr-1_6b')
 model = AutoModelForCausalLM.from_pretrained(
-    'google/gemma-2-2b-it',
-    device_map="auto"
+    'stabilityai/stablelm-2-zephyr-1_6b',
+    device_map="auto",
 )
 
 
@@ -153,7 +153,7 @@ prompt = (
 
 def generate(event):
     # Generar el texto usando el modelo
-    prompt_msg = [{'role': 'user', 'content': prompt+"\n\n"+event}]
+    prompt_msg = [{"role":"system","content":prompt},{'role': 'user', 'content': event}]
     inputs = tokenizer.apply_chat_template(
         prompt_msg,
         add_generation_prompt=False,
@@ -169,4 +169,4 @@ def generate(event):
     
 
     # Imprimir la salida generada
-    return tokenizer.decode(tokens[0], skip_special_tokens=True)
+    return "{".join(tokenizer.decode(tokens[0], skip_special_tokens=True).split("<|user|>")[1].split("{")[1:-1])