BEW / app.py
PatSeal's picture
Update app.py
e130220 verified
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)