Update gen.py

gen.py

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