app.py

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