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

# Load pre-trained model and tokenizer
model_name = "microsoft/DialoGPT-medium"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Streamlit UI
st.title("E-Commerce Chatbot")
st.write("Ask about our products or browse categories. Type 'exit' to end the conversation.")

# Initialize the chat history
chat_history_ids = None

# User input
query = st.text_input("Your Query:")

if query:
    # Tokenize the input and append to chat history
    new_input_ids = tokenizer.encode(query + tokenizer.eos_token, return_tensors="pt")
    
    # Update chat history
    chat_history_ids = new_input_ids if chat_history_ids is None else torch.cat([chat_history_ids, new_input_ids], dim=-1)

    # Generate a response from the model
    response_ids = model.generate(
        chat_history_ids, 
        max_length=1000, 
        pad_token_id=tokenizer.eos_token_id, 
        attention_mask=torch.ones(chat_history_ids.shape, dtype=torch.long)
    )

    # Decode the model's output
    response = tokenizer.decode(response_ids[:, chat_history_ids.shape[-1]:][0], skip_special_tokens=True)
    
    # Display the chatbot's response
    st.write(f"Chatbot: {response}")