app.py

import gradio as gr
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
import os

hf_token = os.getenv("hf_token")


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}, token=hf_token)
tokenizer = AutoTokenizer.from_pretrained(peft_model_id, token=hf_token)

model = PeftModel.from_pretrained(model, peft_model_id)
model.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(
        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:
    {{
    "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>
    """
    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):
    texto_json = generate(pregunta)
    respuesta = ""
    fuente = ""
    tema = ""
    
    if texto_json.startswith(''''{'''):
        try:
            # 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 = [
    ["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"),
        gr.Textbox(label="Fuente"),
        gr.Textbox(label="Tema")
    ],
    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