app.py

import gradio as gr
import json
import re
from huggingface_hub import InferenceClient

client = InferenceClient(
    "mistralai/Mistral-7B-Instruct-v0.1"
)

system_prompt = """[INST] <<SYS>>
You are a helpful, respectful and honest assistant. Always answer in spanish. Never answer in English. Always answer as
helpfully as possible, while being safe.

If a question does not make any sense, or is not factually coherent, explain
why instead of answering something not correct. If you don't know the answer
to a question, please don't share false information.

Your goal is to provide answers relating to a legal document, making classifications of arguments inside an allegation.<</SYS>>
"""

rag_text = """
"""

format_text = """{
    'tipo': "clasificación del argumento (Argumento tipo X)",
    'descripcion': "Descripcion breve del argumento",
    'texto_referencia': "texto original completo perteneciente a dicho argumento"
}"""

prompt_template_text = """
                          --- Lista Argumentos ---
    * Argumento tipo 49 (Desvirtuación de los hechos, el denunciado no está de acuerdo con el relato de los hechos)
    * Argumento tipo 304 (El denunciado solicita medios de prueba)
    * Argumento tipo 994 (El denunciado solicita fotografía como medio de prueba, desafiando el relato del agente denunciante)
    * Argumento tipo 1002 (Prescripción, el delito ha prescrito)
    * Argumento tipo 2014 (Principio de proporcionalidad, denunciado considera sanción excesiva)
    * Argumento tipo 2027 (Niega desobediencia de señales de tráfico de prohibición)
    * Argumento tipo 2002 (Denuncia no notificada en el momento en el que se formuló)
    * Argumento tipo 1001 (Defecto de forma en la denuncia, la sanción contiene algún error)
"""

recorte = """
PRIMERA.- Que los hechos descritos en la denuncia no se corresponden con la realidad de lo
sucedido, por lo que los niego expresamente. La carga de la prueba corresponde al órgano instructor
del procedimiento, en virtud del principio de presunción de inocencia recogido en el art.24 de la
Constitución Española y en el art. 53 de la Ley 39/2015 de 1 de Octubre del Procedimiento
Administrativo Común de las Administraciones Públicas.
"""

def format_prompt(message):
  prompt = "<s>"
  prompt += f"[INST] {message} [/INST]"
  return prompt

def generate(prompt):

    generate_kwargs = dict(
        temperature=0.9,
        max_new_tokens=1024,
        top_p=0.95,
        repetition_penalty=1.0,
        do_sample=True,
        seed=42,
    )

    formatted_prompt = format_prompt(prompt)

    output = client.text_generation(formatted_prompt, **generate_kwargs)

    return output


def process_input(text, rag, prompt_template):

    if not(rag):
        rag_text = " "
    #if not(prompt_template):
    #    prompt_template_text = " "
    
    prompt = f"""
    {system_prompt}

    Teniendo en cuenta que los argumentos se clasifican de la siguiente manera:
    {prompt_template_text}
    
    Partiendo del siguiente fragmento de texto:

    Fragmento:
    ---------------------------------------------------------------
      {text}
    --------------------------------------------------------------

    Identifica y clasifica los argumentos expuestos en el anterior fragmento de texto siguiendo estos pasos:
        Devuelve una lista de JSON, cuya longitud será igual al número de argumentos encontrados, que contenga los argumentos clasificados con el siguiente formato:
        {format_text}


    Recuerda responder en español. Los argumentos que no se puedan clasificar o que *claramente* no encajen con ningún tipo predefinido serán clasificados como 'Desconocido'. Recuerda que los tipos de argumentos de representan con un número que puedes encontrar en la información anterior.
    Respira profundamente y piensa paso a paso.
    [/INST]


    """
    output = generate(prompt)

    # Generamos el archivo JSON
    with open('output2.json', 'w') as f:
        json.dump(output, f)
    
    return output

def create_interface():
    # Definimos los componentes de la interfaz
    input_text = gr.Textbox(label="Input")
    rag_checkbox = gr.Checkbox(label="RAG")
    prompt_template = gr.Checkbox(label="PromptTemplate")
    output_text = gr.Textbox(label="Output")
    classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo")

    # Definimos la función que se ejecutará cuando se envíe la entrada
    def fn(text, rag, prompt_template):
        
        output = process_input(text, rag, prompt_template)
        matches = re.findall(r'"tipo"\s*:\s*"(.*?)"', output)
        classification_types = matches
        classification_types_options = [(option, option) for option in classification_types]
        classification_types_checkboxes = gr.CheckboxGroup(label="Clasificacion Tipo", choices=classification_types_options, interactive = True)
        
        return output, classification_types_checkboxes

    # Definimos la función que se ejecutará cuando se presione el botón de salida
    def mark_checkboxes(output, classification_types_options):
        # Marcamos los checkboxes
        for option in classification_types_options:
            classification_types_checkboxes.set_checkbox_value(optio

    # Definimos la función que se ejecutará cuando se presione el botón de salida
    def save_output(output, classification_types_options):
        # Guardamos los resultados
        with open('output.txt', 'w') as f:
            f.write(output)
        # Marcamos los checkboxes
        for option in classification_types_options:
            classification_types_checkboxes.set_checkbox_value(option[0], True)

n[0], True)
        
    examples = [
        [recorte, False, True]
    ]
        
    # Creamos la interfaz
    iface = gr.Interface(
        fn=fn,
        inputs=[input_text, rag_checkbox, prompt_template],
        outputs=[output_text, classification_types_checkboxes],
        examples=examples
        live=True,
        fn_on_change=mark_checkboxes,
        fn_on_change=save_output
    )

    return iface

iface = create_interface()
iface.launch()