open-assistant / app.py
aidenmp's picture
Update app.py
f52a283 verified
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.")