Update app.py

app.py

@@ -2,12 +2,10 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel
-
 # !python -c "import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'"
 import json
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GemmaTokenizer,  StoppingCriteria, StoppingCriteriaList, GenerationConfig
-# from google.colab import userdata
 import os
 
 #sft_model = "somosnlp/gemma-FULL-RAC-Colombia_v2"
@@ -16,7 +14,6 @@ import os
 sft_model = "somosnlp/RecetasDeLaAbuela_gemma-2b-it-bnb-4bit"
 base_model_name = "unsloth/gemma-2b-it-bnb-4bit"
 
-
 bnb_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_quant_type="nf4",
@@ -42,7 +39,6 @@ model.save_pretrained(".")
 #model.to('cuda')
 tokenizer.save_pretrained(".")
 
-
 class ListOfTokensStoppingCriteria(StoppingCriteria):
     """
     Clase para definir un criterio de parada basado en una lista de tokens específicos.
@@ -71,22 +67,9 @@ stopping_criteria = ListOfTokensStoppingCriteria(tokenizer, stop_tokens)
 stopping_criteria_list = StoppingCriteriaList([stopping_criteria])
 
 def generate_text(prompt, max_length=2100):
-  # prompt="""What were the main contributions of Eratosthenes to the development of mathematics in ancient Greece?"""
   prompt=prompt.replace("\n", "").replace("¿","").replace("?","")
-
-  #EXAMPLE
-  input_text = f'''<bos><start_of_turn>system
-You are a helpful AI assistant.
-Responde en formato json.
-Eres un experto cocinero de la cocina hispanoamericana.<end_of_turn>
-<start_of_turn>user
-¿{prompt}?<end_of_turn>
-<start_of_turn>model
-'''
-
-  inputs = tokenizer.encode(input_text,
-                            return_tensors="pt",
-                            add_special_tokens=False).to("cuda:0")
+  input_text = f'''<bos><start_of_turn>system You are a helpful AI assistant.Eres un experto cocinero de la cocina hispanoamericana.<end_of_turn><start_of_turn>user ¿{prompt}?<end_of_turn><start_of_turn>model'''
+  inputs = tokenizer.encode(input_text, return_tensors="pt", add_special_tokens=False).to("cuda:0")
   max_new_tokens=max_length
   generation_config = GenerationConfig(
                 max_new_tokens=max_new_tokens,
@@ -96,26 +79,15 @@ Eres un experto cocinero de la cocina hispanoamericana.<end_of_turn>
                 repetition_penalty=1.04,  #1.1
                 do_sample=True,
             )
-  outputs = model.generate(generation_config=generation_config,
-                          input_ids=inputs,
-                          stopping_criteria=stopping_criteria_list,)
+  outputs = model.generate(generation_config=generation_config, input_ids=inputs, stopping_criteria=stopping_criteria_list,)
   return tokenizer.decode(outputs[0], skip_special_tokens=False) #True
 
-
-
 def mostrar_respuesta(pregunta):
     try:
       res= generate_text(pregunta, max_length=500)
-      inicio_json = res.find('{')
-      fin_json = res.rfind('}') + 1
-      json_str = res[inicio_json:fin_json]
-      json_obj = json.loads(json_str)
-      # print(json_obj)    
-      return json_obj["Respuesta"]
+      return str(res)
     except Exception as e:
-      json_obj={}
-      json_obj['Respuesta']=str(e)
-      return json_obj
+      return str(e)
 
 # Ejemplos de preguntas
 ejemplos = [
@@ -127,12 +99,10 @@ ejemplos = [
 iface = gr.Interface(
     fn=mostrar_respuesta,
     inputs=gr.Textbox(label="Pregunta"),
-    outputs=[
-        gr.Textbox(label="Respuesta", lines=2),
-    ],
+    outputs=[gr.Textbox(label="Respuesta", lines=2),],
     title="Recetas de la Abuel@",
     description="Introduce tu pregunta sobre recetas de cocina.",
     examples=ejemplos,
 )
 
-iface.queue(max_size=14).launch() # share=True,debug=True
+iface.queue(max_size=14).launch() # share=True,debug=True