File size: 2,191 Bytes
04f76c0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
from transformers import TextDataset, DataCollatorForLanguageModeling
from transformers import Trainer, TrainingArguments
import torch
# Define your model and tokenizer
model_name = "gpt2" # You can use different GPT-2 variants like "gpt2-medium" if needed
config = GPT2Config.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name, config=config)
# Move the model to the appropriate device (e.g., GPU if available)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
# Define training arguments
training_args = TrainingArguments(
output_dir="./fitnessbot", # Directory to save model checkpoints and results
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=2,
save_steps=1000,
save_total_limit=2,
)
# Load your dataset
dataset = TextDataset(
tokenizer=tokenizer,
file_path="fitness_data.txt", # Path to your dataset file
block_size=256, # Adjust this block size based on your model's token limit
)
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False,
)
# Create a Trainer instance and fine-tune the model
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=dataset,
)
# Fine-tune the model
trainer.train()
trainer.save_model()
trainer.save_state()
# Chat with the bot
print("Welcome to the Fitness Chatbot. Type 'exit' to end the conversation.")
while True:
user_input = input("You: ")
if user_input.lower() == 'exit':
print("Fitness Chatbot: Goodbye!")
break
# Encode the new user input, add to the chat history, and generate a response
input_ids = tokenizer.encode("You: " + user_input, return_tensors="pt").to(device)
bot_input_ids = input_ids
# Generate a response
response = model.generate(bot_input_ids, max_length=100, num_return_sequences=1, no_repeat_ngram_size=2)
bot_response = tokenizer.decode(response[0], skip_special_tokens=True)
print("Fitness Chatbot:", bot_response) |