Update app.py

app.py

@@ -2,43 +2,21 @@ 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
 import os
 
-#sft_model = "somosnlp/gemma-FULL-RAC-Colombia_v2"
 #sft_model = "somosnlp/RecetasDeLaAbuela_mistral-7b-instruct-v0.2-bnb-4bit"
 #base_model_name = "unsloth/Mistral-7B-Instruct-v0.2"
-sft_model1 = "somosnlp/RecetasDeLaAbuela_gemma-2b-it-bnb-4bit"
-sft_model2 = "somosnlp/RecetasDeLaAbuela_mistral-7b-instruct-v0.2-bnb-4bit"
+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",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
 max_seq_length=400
-
-# if torch.cuda.get_device_capability()[0] >= 8:
-#     # print("Flash Attention")
-#     attn_implementation="flash_attention_2"
-# else:
-#     attn_implementation=None
-attn_implementation=None
-
-#base_model = AutoModelForCausalLM.from_pretrained(model_name,return_dict=True,torch_dtype=torch.float16,)
 base_model = AutoModelForCausalLM.from_pretrained(base_model_name,return_dict=True,device_map="auto", torch_dtype=torch.float16,)
-#base_model = AutoModelForCausalLM.from_pretrained(base_model_name, return_dict=True, device_map = {"":0}, attn_implementation = attn_implementation,).eval()
-
 tokenizer = AutoTokenizer.from_pretrained(base_model_name, max_length = max_seq_length)
-sft_model = sft_model1
 ft_model = PeftModel.from_pretrained(base_model, sft_model)
 model = ft_model.merge_and_unload()
 model.save_pretrained(".")
-#model.to('cuda')
 tokenizer.save_pretrained(".")
 
 class ListOfTokensStoppingCriteria(StoppingCriteria):
@@ -68,15 +46,9 @@ stopping_criteria = ListOfTokensStoppingCriteria(tokenizer, stop_tokens)
 # Añade tu criterio de parada a una StoppingCriteriaList
 stopping_criteria_list = StoppingCriteriaList([stopping_criteria])
 
-def generate_text(modelin, prompt, context, max_length=2100):
-  print('Modelo es: '+modelin)
-  #sft_model = modelin
-  #ft_model = PeftModel.from_pretrained(base_model, sft_model)
-  #model = ft_model.merge_and_unload()
-
+def generate_text(prompt, context, max_length=2100):
   prompt=prompt.replace("\n", "").replace("¿","").replace("?","")
-  #input_text = f'''<bos><start_of_turn>system ¿{context}?<end_of_turn><start_of_turn>user ¿{prompt}?<end_of_turn><start_of_turn>model'''
-  input_text = str(context)+str(prompt)
+  input_text = f'''<bos><start_of_turn>system ¿{context}?<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(
@@ -90,25 +62,23 @@ def generate_text(modelin, prompt, context, max_length=2100):
   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(modelo, pregunta, contexto):
+def mostrar_respuesta(pregunta, contexto):
     try:
-      print('Modelo: '+str(modelo))
-      print('Pregunta: '+str(pregunta))
-      print('Contexto: '+str(contexto))
-      res= generate_text(modelo, pregunta, contexto, max_length=500)    
-      print('Respuesta: '+str(contexto))
+      res= generate_text(pregunta, contexto, max_length=500)
       return str(res)
     except Exception as e:
       return str(e)
 
 # Ejemplos de preguntas
-mis_ejemplos = [[sft_model1,"¿Dime la receta de la tortilla de patatatas?","Cocinero español"],[sft_model1,"¿Dime la receta del ceviche?","Cocinero peruano"], [sft_model1, "¿Como se cocinan unos autenticos frijoles?","Cocinero de México"],]
-lista_modelos = [sft_model1, sft_model2]
+mis_ejemplos = [
+    ["¿Dime la receta de la tortilla de patatatas?"],
+    ["¿Dime la receta del ceviche?"],
+    ["¿Como se cocinan unos autenticos frijoles?"],
+]
 
 iface = gr.Interface(
     fn=mostrar_respuesta,
-    inputs=[gr.Dropdown(choices=lista_modelos, value = sft_model1, label="Modelo", type="value"), gr.Textbox(label="Pregunta"), 
-        gr.Textbox(label="Contexto", value="You are a helpful AI assistant. Eres un experto cocinero hispanoamericano."),],
+    inputs=[gr.Textbox(label="Pregunta"), gr.Textbox(label="Contexto", value="You are a helpful AI assistant. Eres un experto cocinero hispanoamericano."),],
     outputs=[gr.Textbox(label="Respuesta", lines=2),],
     title="Recetas de la Abuel@",
     description="Introduce tu pregunta sobre recetas de cocina.",