Spaces:
Sleeping
Sleeping
| 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") |