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

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

# Function to generate a response
def generate_response(chat_history_ids, new_user_input):
    new_user_input_ids = tokenizer.encode(new_user_input + tokenizer.eos_token, return_tensors='pt')
    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if chat_history_ids is not None else new_user_input_ids
    return model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)

# Streamlit app
st.title("DialoGPT Chat")

# Initialize chat history
chat_history_ids = None

# Conversation loop
while st.button("Restart Conversation"):
    # Get user input
    user_input = st.text_input("You:", "")

    if user_input:
        # Generate response
        chat_history_ids = generate_response(chat_history_ids, user_input)

        # Display DialoGPT response
        st.text("DialoGPT: {}".format(tokenizer.decode(chat_history_ids[:, -1][0], skip_special_tokens=True)))

# Inform the user that the conversation has ended
st.text("Conversation Ended. Press the 'Restart Conversation' button to start a new conversation.")