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 = """
    SEGUNDA.- Que como consecuencia de lo manifestado con anterioridad, intereso la apertura del
    periodo de prueba, conforme al el art. 13 del Real Decreto 320/1994 en relación con los arts. 77 y 78
    de la Ley 39/2015 y propongo la práctica de los medios de prueba, conforme a lo dispuesto en el 13.1
    del RD. 320/1994, que esta parte considera imprescindibles para poder ejercer su derecho a la
    defensa, con expresa petición de la notificación de su resultado y sin perjuicio de que una vez recibidas
    pueda realizar cuantas alegaciones estime pertinentes, que se deja interesada como sigue:

    A- Documental, para que se incorpore al expediente INFORME DEL AGENTE DENUNCIANTE
    relativo a las circunstancias de la presunta infracción, con expresión en concreto de su consideración
    acerca de la relación entre la conducta del presunto infractor y su determinación como infracción
    tipificada
"""

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)
    
    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,value=[option[0] for option in classification_types_options], interactive = True)

        return output, classification_types_checkboxes
        
    examples = [
        [recorte, False, True]
    ]

    HF_TOKEN = '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'
    hf_writer = gr.HuggingFaceDatasetSaver(HF_TOKEN, "MistralAI")

    # Creamos la interfaz
    iface = gr.Interface(
        fn=fn,
        inputs=[input_text, rag_checkbox, prompt_template],
        outputs=[output_text, classification_types_checkboxes],
        examples=examples,
        allow_flagging="manual",
        flagging_options=["👍", "👎"],
        flagging_callback=hf_writer
    )

    return iface

iface = create_interface()
iface.launch()