Update app.py

app.py

@@ -1,34 +1,20 @@
 import gradio as gr
 
-# !python -c "import torch; assert torch.cuda.get_device_capability()[0] >= 8, 'Hardware not supported for Flash Attention'"
-import json
 import torch
+from peft import PeftModel, PeftConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, AutoTokenizer, GenerationConfig
+
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, GemmaTokenizer,  StoppingCriteria, StoppingCriteriaList, GenerationConfig
-# from google.colab import userdata
-import os
 
-model_id = "somosnlp/kuntur-peru-legal-es-gemma-2b-it-merged"
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,
-    bnb_4bit_quant_type="nf4",
-    bnb_4bit_compute_dtype=torch.bfloat16
-)
-max_seq_length=512
+import os
 
-# if torch.cuda.get_device_capability()[0] >= 8:
-#     # print("Flash Attention")
-#     attn_implementation="flash_attention_2"
-# else:
-#     attn_implementation=None
-attn_implementation=None
+peft_model_id = "daqc/kuntur-peru-legal-es-gemma-2b-it"
+config = PeftConfig.from_pretrained(peft_model_id)
+model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path, return_dict=True, load_in_8bit=True, device_map={"":0})
+tokenizer = AutoTokenizer.from_pretrained(peft_model_id)
 
-tokenizer = AutoTokenizer.from_pretrained(model_id,
-                                          max_length = max_seq_length)
-model = AutoModelForCausalLM.from_pretrained(model_id,
-                                            #  quantization_config=bnb_config,
-                                             device_map = {"":0},
-                                              attn_implementation = attn_implementation, # A100 o H100
-                                             ).eval()
+model = PeftModel.from_pretrained(model, peft_model_id)
+model.eval()
 
 
 
@@ -51,7 +37,7 @@ class ListOfTokensStoppingCriteria(StoppingCriteria):
         return False
 
 # Uso del criterio de parada personalizado
-stop_tokens = ["end_of_turn"]  # Lista de tokens de parada
+stop_tokens = ["<end_of_turn>"]  # Lista de tokens de parada
 
 # Inicializa tu criterio de parada con el tokenizer y la lista de tokens de parada
 stopping_criteria = ListOfTokensStoppingCriteria(tokenizer, stop_tokens)
@@ -59,13 +45,17 @@ stopping_criteria = ListOfTokensStoppingCriteria(tokenizer, stop_tokens)
 # Añade tu criterio de parada a una StoppingCriteriaList
 stopping_criteria_list = StoppingCriteriaList([stopping_criteria])
 
-def generate_text(prompt, max_length=2048):
-  # 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'''<start_of_turn>system
+def generate(
+        instruction,
+        max_new_tokens=256,
+        temperature=0.1,
+        top_p=0.75,
+        top_k=40,
+        num_beams=2,
+        **kwargs,
+):
+    instruction = instruction.replace("¿","").replace("?","")
+    system = f"""<start_of_turn>system
     You are a helpful AI assistant. You only answer in JSON format.
     Eres un agente experto en la constitución política del perú de 1993 que solo responde formato JSON:
     {{
@@ -79,54 +69,69 @@ def generate_text(prompt, max_length=2048):
     3. tema: Solo escoge los temas de la lista proporcionada, no inventes ni crees un nuevo tema, en caso de considerarse mas de 2 temas se separa con punto y coma, escoge solo los que se adecuen a la respuesta, no consideres todos los temas al mismo tiempo: Educación, Conflictos sociales, Prevención de la corrupción, Servicios públicos, Violencia contra la niñez, Desigualdad y violencia hacia las mujeres, Seguridad ciudadana, Discapacidad o Salud.
     UNICAMENTE DEBES RESPONDER EN FORMATO JSON, SOLO EN JSON, JSON, JSON
     <end_of_turn>
-    <start_of_turn>user
-    ¿{prompt}?<end_of_turn>\n<start_of_turn>model\n'''
-
-  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,
-                temperature=0.15,
-                top_p=0.75,    #0.9,
-                top_k=40, # 45
-                num_beams=2, #me
-                repetition_penalty=1.,  #1.1
-                do_sample=True,
-            )
-  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
+    """
+    prompt = f"""{system} <start_of_turn>user
+    {instruction}<end_of_turn> <start_of_turn>model\n
+    """
+    print(prompt)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    input_ids = inputs["input_ids"].to("cuda")
+    attention_mask = inputs["attention_mask"].to("cuda")
+    generation_config = GenerationConfig(
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        num_beams=num_beams,
+        **kwargs,
+    )
+    with torch.no_grad():
+        generation_output = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            generation_config=generation_config,
+            return_dict_in_generate=True,
+            #output_scores=True,
+            max_new_tokens=max_new_tokens,
+            early_stopping=True
+        )
+    s = generation_output.sequences[0]
+    output = tokenizer.decode(s, skip_special_tokens=True)
+    return output.split("model")[1]
+
+
 
 
 
 def mostrar_respuesta(pregunta):
-    respuesta_default = "No se pudo generar una respuesta adecuada."
-    json_obj = {
-        "respuesta": respuesta_default,
-        "fuente": respuesta_default,
-        "tema": respuesta_default
-    }
+    texto_json = generate(pregunta)
+    respuesta = ""
+    fuente = ""
+    tema = ""
     
-    if pregunta:
+    if texto_json.startswith('{'):
         try:
-            res = generate_text(pregunta, max_length=512)
-            inicio_json = res.find('{')
-            fin_json = res.rfind('}') + 1
-            json_str = res[inicio_json:fin_json]
-            json_obj = json.loads(json_str)
-
-            # Verificar si el JSON contiene todas las claves necesarias
-            if all(key in json_obj for key in ["respuesta", "fuente", "tema"]):
-                return json_obj["respuesta"], json_obj["fuente"], json_obj["tema"]
-            else:
-                return res, respuesta_default, respuesta_default
-        except Exception as e:
-            print("Error al procesar la respuesta:", e)
-    
-    return json_obj["respuesta"], json_obj["fuente"], json_obj["tema"]
+            # Busca las posiciones de inicio y fin de cada campo
+            inicio_respuesta = texto_json.find('"respuesta":') + len('"respuesta":')
+            fin_respuesta = texto_json.find('"fuente":')
+            inicio_fuente = texto_json.find('"fuente":') + len('"fuente":')
+            fin_fuente = texto_json.find('"tema":')
+            inicio_tema = texto_json.find('"tema":') + len('"tema":')
+
+            # Extrae los valores de cada campo
+            respuesta = texto_json[inicio_respuesta:fin_respuesta].strip().strip('"')
+            fuente = texto_json[inicio_fuente:fin_fuente].strip().strip('"')
+            # Verifica si la clave "tema" existe en el JSON antes de extraer su valor
+            if '"tema":' in texto_json:
+                tema = texto_json[inicio_tema:].strip().strip('"')
+        except ValueError:
+            pass
+    else:
+        respuesta = texto_json.strip().strip('"')
+
+    return respuesta, fuente, tema
+
+
+
 
 # Ejemplos de preguntas
 ejemplos = [
@@ -143,7 +148,7 @@ iface = gr.Interface(
         gr.Textbox(label="Fuente", lines=1),
         gr.Textbox(label="Tema", lines=1)
     ],
-    title="HolaaaaaaaaaaaaaaaaaConsulta Juridica basada en la Constitucion Politica del Peru",
+    title="Consulta Juridica basada en la Constitucion Politica del Peru",
     description="Introduce tu pregunta sobre la Constitución Politica o una situación donde creas que tus derechos hayan sido vulnerados.",
     examples=ejemplos,
 )