import gradio as gr
import json
from huggingface_hub import InferenceClient

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

rag_text = "Este es el texto RAG"
prompt_template_text = "Este es el texto del template de prompt"

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)

    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
    output = ""

    for response in stream:
        output += response.token.text
        yield output
    return output


def process_input(text, rag, prompt_template):
    prompt = text
    if rag:
        prompt += rag_text
    if prompt_template:
        prompt += prompt_template_text
    json_text = generate(prompt)
    
    # Convertimos el generador en una cadena JSON
    #json_str = ''.join(json_text)

    # Convertimos la cadena JSON en un objeto JSON
    #json_obj = json.loads(json_str)

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

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)
        with open('output.json', 'r') as f:
            data = json.load(f)
        classification_types = [item['clasificacion_tipo'] for item in data]
        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

    # Creamos la interfaz
    iface = gr.Interface(
        fn=fn,
        inputs=[input_text, rag_checkbox, prompt_template],
        outputs=[output_text, classification_types_checkboxes]
    )

    return iface

iface = create_interface()
iface.launch()