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

# Define the training arguments
training_args = TrainingArguments(
    per_device_train_batch_size=4,
    num_train_epochs=3,
    logging_dir='./logs',
)

# Load the pre-trained GPT-2 model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = GPT2LMHeadModel.from_pretrained("gpt2")

# Example training data
training_data = [
    "What is your name?",
    "How are you?",
    "What do you do?",
    "Tell me about yourself."
]

# Tokenize the training data
input_ids = tokenizer(training_data, return_tensors="pt", padding=True, truncation=True)["input_ids"]

# Define a dummy data collator (required by Trainer)
class DummyDataCollator:
    def __call__(self, features):
        return features

# Define a Trainer instance
trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=DummyDataCollator(),
    train_dataset=input_ids
)

# Train the model
trainer.train()

# Define the chatbot function
def chatbot(input_text):
    # Tokenize input text
    input_ids = tokenizer.encode(input_text, return_tensors="pt")

    # Generate response from the model
    output_ids = model.generate(input_ids, max_length=50, pad_token_id=tokenizer.eos_token_id)
    
    # Decode the generated response
    response = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    return response

# Create the Gradio interface
chatbot_interface = gr.Interface(chatbot, "textbox", "textbox", title="Chatbot")

# Launch the Gradio interface
chatbot_interface.launch()