Upload app.py

app.py

@@ -0,0 +1,144 @@
+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 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
+
+# if torch.cuda.get_device_capability()[0] >= 8:
+#     # print("Flash Attention")
+#     attn_implementation="flash_attention_2"
+# else:
+#     attn_implementation=None
+attn_implementation=None
+
+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()
+
+
+
+class ListOfTokensStoppingCriteria(StoppingCriteria):
+    """
+    Clase para definir un criterio de parada basado en una lista de tokens específicos.
+    """
+    def __init__(self, tokenizer, stop_tokens):
+        self.tokenizer = tokenizer
+        # Codifica cada token de parada y guarda sus IDs en una lista
+        self.stop_token_ids_list = [tokenizer.encode(stop_token, add_special_tokens=False) for stop_token in stop_tokens]
+
+    def __call__(self, input_ids, scores, **kwargs):
+        # Verifica si los últimos tokens generados coinciden con alguno de los conjuntos de tokens de parada
+        for stop_token_ids in self.stop_token_ids_list:
+            len_stop_tokens = len(stop_token_ids)
+            if len(input_ids[0]) >= len_stop_tokens:
+                if input_ids[0, -len_stop_tokens:].tolist() == stop_token_ids:
+                    return True
+        return False
+
+# Uso del criterio de parada personalizado
+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)
+
+# 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
+    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:
+    {{
+    "respuesta": "La defensa de la persona humana y el respeto de su dignidad son el fin supremo de la sociedad y del Estado.", 
+    "fuente": "Constitución Política del Perú 1993 - Capitulo 1, Artículo 1;",
+    "tema": "Violencia contra la niñez; Desigualdad y violencia hacia las mujeres; Seguridad ciudadana"
+    }}
+    Respondes nada mas que en formato JSON y respeta las instrucciones para las siguientes claves, IMPORTANTE NO CREAR SUBCLAVES, toda la respuesta debe ir dentro de la clave 'respuesta':
+    1. respuesta: Solo en esta seccion se detallara toda, TODA, TODA la respuesta, esto incluye la resolución del problema propuesto o la solucion a situaciones de vulneración de derechos.
+    2. fuente: La fuente debe ser siempre 'Constitución Política del Perú 1993', variando únicamente el número de artículo y capítulo. Solamente debes basarte en los articulos de la Constitucion Politica del Peru de 1993, puedes usar otro texto juridico solo en casos extremos, en este caso se debera separar los textos juridicos con sus respectivos capitulos y articulos con punto y coma.
+    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
+
+
+
+def mostrar_respuesta(pregunta):
+    json_obj={}
+    json_obj['respuesta']='Esperando'
+    json_obj['fuente']='Esperando'
+    json_obj['tema']='Esperando'
+    if pregunta!="":
+        try:
+          res= generate_text(pregunta, max_length=512)
+          # print(">> RES:",res)
+          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:",json_obj)    
+          return json_obj["respuesta"], json_obj["fuente"], json_obj["tema"]
+        except:          
+          return json_obj["respuesta"], json_obj["fuente"], json_obj["tema"]    
+    return json_obj["respuesta"], json_obj["fuente"], json_obj["tema"]
+
+# Ejemplos de preguntas
+ejemplos = [
+    ["Una persona ha sido despedida de su trabajo injustamente y necesita entender cuáles son sus derechos laborales según la Constitución Política del Perú."],
+    ["¿Que establece el articulo 2 de la Constitucion Politica del Peru respecto a los derechos humanos?"],
+    ["¿Que articulo de la Constitucion Politica del Peru establece que toda persona tiene derecho a la vida, a su identidad, a su integridad moral, psíquica y física y a su libre desarrollo y bienestar?"],
+]
+
+iface = gr.Interface(
+    fn=mostrar_respuesta,
+    inputs=gr.Textbox(label="Pregunta"),
+    outputs=[
+        gr.Textbox(label="Respuesta", lines=2),
+        gr.Textbox(label="Fuente", lines=1),
+        gr.Textbox(label="Tema", lines=1)
+    ],
+    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,
+)
+
+iface.queue(max_size=14).launch(debug=True) # share=True,debug=True