import streamlit as st
import os
from transformers import AutoModelForCausalLM, AutoTokenizer

# Obtener la API key desde la variable de entorno
api_key = os.getenv("HF_API_KEY")

# Cargar el modelo de OpenAssistant con la API key
model = AutoModelForCausalLM.from_pretrained("OpenAssistant/falcon-7b-sft-mix-2000", use_auth_token=api_key, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("OpenAssistant/falcon-7b-sft-mix-2000", use_auth_token=api_key)

# Función para generar una respuesta usando el modelo cargado
def generate_response(prompt):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Crear la interfaz de usuario con Streamlit
st.title("Asistente de OpenAssistant")
st.write("Introduce un texto para obtener una respuesta generada:")

user_input = st.text_area("Tu mensaje:", height=150)

if st.button("Generar respuesta"):
    if user_input:
        with st.spinner("Generando respuesta..."):
            response = generate_response(user_input)
            st.subheader("Respuesta generada:")
            st.write(response)
    else:
        st.warning("Por favor, ingresa un mensaje.")