|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments |
|
|
|
|
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium") |
|
|
model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-medium") |
|
|
|
|
|
|
|
|
dataset = [ |
|
|
("Hello!", "Hi there!"), |
|
|
("How are you?", "I'm doing well, thanks!"), |
|
|
|
|
|
] |
|
|
|
|
|
|
|
|
tokenized_dataset = tokenizer([example[0] for example in dataset], return_tensors="pt", padding=True, truncation=True) |
|
|
|
|
|
|
|
|
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() |
|
|
|
|
|
|
|
|
model.save_pretrained("fine_tuned_dialogpt") |
|
|
|
|
|
|
|
|
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 |
|
|
|
|
|
|
|
|
while True: |
|
|
user_input = input("You: ") |
|
|
if user_input.lower() == 'exit': |
|
|
print("Chatbot: Goodbye!") |
|
|
break |
|
|
response = chatbot(user_input) |
|
|
print("Chatbot:", response) |
|
|
|