import gradio as gr
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

model_name = "mrm8488/gpt2-finetuned-jhegarty-texts"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)

tokenizer.pad_token = tokenizer.eos_token
model.config.pad_token_id = model.config.eos_token_id

def chat_with_gpt2(user_input):
    prompt = f"You are a helpful travel assistant. Answer concisely and informatively. Question: {user_input}\nAnswer:"
    inputs = tokenizer.encode(prompt, return_tensors="pt", truncation=True, max_length=512)
    outputs = model.generate(
        inputs,
        max_length=250,
        num_return_sequences=1,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        top_p=0.95,
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(prompt):].strip()

interface = gr.Interface(
    fn=chat_with_gpt2,
    inputs=gr.Textbox(lines=3, placeholder="Ask about your trip..."),
    outputs="text",
    title="Trip Planner Chatbot",
    description="Ask anything related to your travel plans!"
)

interface.launch()