import streamlit as st
import requests
from huggingface_hub import HfApi
import transformers
from transformers import AutoModelForSeq2SeqLM

# Define la función para procesar el documento con el prompt
def process_document(document, prompt, model_name):
    # Carga el modelo
    model = transformers.AutoModelForSeq2SeqLM.from_pretrained(model_name)

    # Carga el tokenizer
    tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)

    # Tokeniza el texto del documento y del prompt
    input_ids = tokenizer.encode(document + prompt, return_tensors="pt")

    # Genera la respuesta del modelo
    response = model.generate(input_ids, max_length=512)

    # Decodifica la respuesta en texto
    model_response = tokenizer.decode(response[0], skip_special_tokens=True)

    # Devuelve la respuesta del modelo
    return model_response

# Crea la interfaz de usuario de Streamlit
st.title("Aplicación de Procesamiento de Documentos con Prompt")

# Campo de texto para el documento
document_input = st.text_area("Ingrese el documento:")

# Campo de texto para el prompt
prompt_input = st.text_area("Ingrese el prompt:")

# Selector de modelo
model_options = ["v2ray/Mixtral-8x22B-v0.1", "facebook/bart-large-cnn", ...]
selected_model = st.selectbox("Seleccione un modelo:", model_options)

# Botón para procesar el documento
if st.button("Procesar"):
    # Procesa el documento con el prompt y el modelo seleccionado
    model_response = process_document(document_input, prompt_input, selected_model)

    # Muestra la respuesta del modelo
    st.write("**Respuesta del modelo:**")
    st.write(model_response)

# Define la configuración de la aplicación
app_config = {
    "title": "Aplicación de Procesamiento de Documentos con Prompt",
    "description": "Esta aplicación permite ingresar un documento, un prompt y enviarlos a un modelo para obtener una respuesta.",
    "keywords": ["procesamiento de lenguaje natural", "prompt", "modelo"],
    "category": "apps",
}

# Despliega la aplicación en Hugging Face Spaces (descomentar para desplegar)
# push_to_hub(app=app, config=app_config)