import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer

# Load the model and tokenizer from Hugging Face model hub
model_name = "gpt2"  # You can replace "gpt2" with your fine-tuned model if you have one
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

def chat_with_me(input_text, history=[]):
    # Encode the new input with history
    new_input_ids = tokenizer.encode(input_text + tokenizer.eos_token, return_tensors="pt")

    # Append the new user input to the chat history
    bot_input_ids = torch.cat([torch.tensor(history), new_input_ids], dim=-1) if history else new_input_ids

    # Generate the model's response
    history_ids = model.generate(bot_input_ids, max_length=1000, pad_token_id=tokenizer.eos_token_id)

    # Decode the response and append to history
    response = tokenizer.decode(history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
    history += new_input_ids.tolist()

    return response, history

with gr.Blocks() as demo:
    with gr.Row():
        chat = gr.Chatbot()
        user_input = gr.Textbox(placeholder="Ask me anything...")
    with gr.Row():
        clear = gr.Button("Clear")

    # Define interactions
    user_input.submit(chat_with_me, [user_input, chat], [chat, user_input])
    clear.click(lambda: None, None, chat)

demo.launch()