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

# Load the fine-tuned DialoGPT model and tokenizer from Hugging Face
@st.cache_resource  # Cache the model to avoid reloading every time
def load_model():
    try:
        # Load the fine-tuned model and tokenizer from the Hugging Face Hub
        st.write("Loading model and tokenizer...")
        model = AutoModelForCausalLM.from_pretrained("username/fine-tuned-dialoGPT-crm-chatbot")
        tokenizer = AutoTokenizer.from_pretrained("username/fine-tuned-dialoGPT-crm-chatbot")
        return model, tokenizer
    except Exception as e:
        st.error(f"Failed to load the model or tokenizer: {e}")
        return None, None

# Define the chatbot function that generates responses
def generate_response(model, tokenizer, input_text, max_length=100):
    try:
        # Tokenize the input text
        inputs = tokenizer.encode(input_text, return_tensors="pt")
        
        # Generate a response using the model
        outputs = model.generate(inputs, max_length=max_length, pad_token_id=tokenizer.eos_token_id)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        return response
    except Exception as e:
        st.error(f"Error during text generation: {e}")
        return "Sorry, something went wrong while generating the response."

# Main Streamlit app function
def chatbot_app():
    # Load model and tokenizer
    model, tokenizer = load_model()
    
    if model is None or tokenizer is None:
        st.error("Unable to load the chatbot model. Please check API or model availability.")
        return
    
    st.title("CRM Chatbot")
    st.write("This chatbot helps with customer service inquiries. Feel free to ask anything!")

    # Chat history to maintain a conversation flow
    if 'chat_history' not in st.session_state:
        st.session_state['chat_history'] = []

    # Input box for user message
    user_input = st.text_input("You:", value="", key="input")
    
    # Submit button to generate a response
    if st.button("Send") and user_input:
        # Display the user's message
        st.session_state.chat_history.append(f"You: {user_input}")

        # Generate chatbot's response
        bot_response = generate_response(model, tokenizer, user_input)
        st.session_state.chat_history.append(f"Chatbot: {bot_response}")

    # Display the chat history
    for message in st.session_state.chat_history:
        st.write(message)

# Run the chatbot app
if __name__ == "__main__":
    chatbot_app()