import streamlit as st
from transformers import pipeline

# Cargar el modelo de DeepSeek AI para el chatbot
@st.cache_resource
def load_model():
    model_name = "deepseek-ai/DeepSeek-R1"  # Cambia esto si el modelo es diferente
    
    # Usar la tarea 'text-generation' en lugar de 'conversational'
    chatbot = pipeline("text-generation", model=model_name, trust_remote_code=True)
    return chatbot

# Inicializa el modelo
chatbot = load_model()

def main():
    # Título de la aplicación
    st.title("ChatBot - DeepSeek AI")

    # Introducción
    st.write("¡Bienvenido al ChatBot! Puedes hacerle preguntas y recibir respuestas.")

    # Campo de entrada de texto
    user_input = st.text_input("Escribe tu pregunta:", "")

    if user_input:
        # Generar una respuesta del chatbot usando el modelo
        response = chatbot(user_input, max_length=100, num_return_sequences=1)
        
        # Mostrar la respuesta generada
        st.write(f"**Respuesta del chatbot:** {response[0]['generated_text']}")

    # Botón para reiniciar el chat (opcional)
    if st.button("Reiniciar chat"):
        st.write("Chat reiniciado. Puedes comenzar de nuevo.")

if __name__ == "__main__":
    main()