File size: 2,130 Bytes
c6a61ed |
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 |
import torch
from transformers import GPT2Tokenizer, GPT2LMHeadModel, AdamW, get_linear_schedule_with_warmup,
from datasets import load_dataset
# Load the jokes dataset
dataset = load_dataset("ysharma/short_jokes")
# Accessing the train split
train_data = dataset['train']
# Shuffle the dataset and select 20% of the data
twenty_percent_size = int(0.2 * len(train_data))
subset = train_data.shuffle(seed=42)[:twenty_percent_size]
# Use GPT-2's tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
tokenizer.pad_token = tokenizer.eos_token
# Tokenize the dataset
def tokenize_function(examples):
return tokenizer(examples["Joke"], padding="max_length", truncation=True, max_length=50)
tokenized_datasets = dataset.map(tokenize_function, batched=True)
# Load GPT-2 model
model = GPT2LMHeadModel.from_pretrained("gpt2-medium")
model.train()
# Training parameters
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
optimizer = AdamW(model.parameters(), lr=5e-5)
num_epochs = 100
total_steps = len(tokenized_datasets["train"]) * num_epochs
scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=0, num_training_steps=total_steps)
# Training loop
for epoch in range(num_epochs):
for idx, batch in enumerate(tokenized_datasets["train"]):
inputs = torch.tensor(batch["input_ids"]).to(device)
attention_mask = torch.tensor(batch["attention_mask"]).to(device)
outputs = model(inputs, attention_mask=attention_mask, labels=inputs)
loss = outputs.loss
loss.backward()
optimizer.step()
scheduler.step()
optimizer.zero_grad()
if idx % 100 == 0:
print(f"Epoch: {epoch}, Batch: {idx}, Loss: {loss.item()}")
if epoch%5==0:
save_directory = f"./trained_gpt2_jokes/{epoch}"
model.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)
print("Training completed!")
save_directory = "./trained_gpt2_jokes/final"
model.save_pretrained(save_directory)
tokenizer.save_pretrained(save_directory)
print(f"Model and tokenizer saved to {save_directory}") |