Spaces:
Sleeping
Sleeping
File size: 1,756 Bytes
2da2440 2638e3e 2da2440 2638e3e 2da2440 2638e3e 2da2440 2638e3e | 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 | import torch
from torch.utils.data import TensorDataset, DataLoader
from transformers import GPT2Config, GPT2LMHeadModel
from torch.optim import AdamW
data = torch.load("train_data.pt")
dataset = TensorDataset(
data["input_ids"],
data["attention_mask"],
data["labels"]
)
batch_size = 8
loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
if torch.backends.mps.is_available():
device = torch.device("mps")
elif torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
model = GPT2LMHeadModel.from_pretrained("gpt2")
model.to(device)
optimizer = AdamW(model.parameters(), lr=5e-5)
epochs = 3
print("\n--- starting training... ---")
for epoch in range(epochs):
model.train()
total_loss = 0.0
for step, batch in enumerate(loader):
b_input_ids = batch[0].to(device)
b_attention_mask = batch[1].to(device)
b_labels = batch[2].to(device)
optimizer.zero_grad()
outputs = model(
input_ids=b_input_ids,
attention_mask=b_attention_mask,
labels=b_labels
)
loss = outputs.loss
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
optimizer.step()
total_loss += loss.item()
if step % 100 == 0 and step > 0:
avg_step_loss = total_loss / step
print(f"Epoch {epoch + 1}/{epochs} | Step {step}/{len(loader)} | Current Loss: {loss.item():.4f} | Running Avg: {avg_step_loss:.4f}")
epoch_avg = total_loss / len(loader)
print(f"\n>>> Epoch {epoch + 1} Complete | Average Loss: {epoch_avg:.4f} <<<\n")
torch.save(model.state_dict(), "trained_weights.pt")
print("training complete.\n") |