from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments

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

# Prepare your dataset (replace this with your own dataset)
dataset = [
    ("Hello!", "Hi there!"),
    ("How are you?", "I'm doing well, thanks!"),
    # Add more conversational pairs here
]

# Tokenize the dataset
tokenized_dataset = tokenizer([example[0] for example in dataset], return_tensors="pt", padding=True, truncation=True)

# Fine-tune the model
training_args = TrainingArguments(
    per_device_train_batch_size=4,
    num_train_epochs=3,
    logging_dir='./logs',
    logging_steps=100,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
)

trainer.train()

# Save the fine-tuned model
model.save_pretrained("fine_tuned_dialogpt")

# Example of using the fine-tuned model for chatbot
def chatbot(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    response_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
    response_text = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
    return response_text

# Example interaction with the fine-tuned model
while True:
    user_input = input("You: ")
    if user_input.lower() == 'exit':
        print("Chatbot: Goodbye!")
        break
    response = chatbot(user_input)
    print("Chatbot:", response)