File size: 4,346 Bytes
7b41c88 |
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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 |
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score
from datasets import Dataset
import pytorch_lightning as pl
from transformers import RobertaTokenizer, RobertaModel
from pytorch_lightning.callbacks import ModelCheckpoint, EarlyStopping
from pytorch_lightning.strategies import DDPStrategy
class RoBERTaClassifier(pl.LightningModule):
def __init__(self, num_labels=4, lr=2e-5, class_weights=None):
super().__init__()
self.save_hyperparameters()
self.model = RobertaModel.from_pretrained("roberta-base", add_pooling_layer=False)
self.dropout = nn.Dropout(0.3)
self.classifier = nn.Linear(self.model.config.hidden_size, num_labels)
if class_weights is not None:
weights = torch.tensor(class_weights, dtype=torch.float32)
self.loss_fn = nn.CrossEntropyLoss(weight=weights)
else:
self.loss_fn = nn.CrossEntropyLoss()
def forward(self, input_ids, attention_mask):
outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
cls_output = outputs.last_hidden_state[:, 0, :]
cls_output = self.dropout(cls_output)
return self.classifier(cls_output)
def training_step(self, batch, batch_idx):
input_ids, attention_mask, labels = batch["input_ids"], batch["attention_mask"], batch["label"]
logits = self(input_ids, attention_mask)
loss = self.loss_fn(logits, labels)
preds = torch.argmax(logits, dim=1)
acc = accuracy_score(labels.cpu(), preds.cpu())
self.log("train_loss", loss, prog_bar=True)
self.log("train_acc", acc, prog_bar=True)
return loss
def validation_step(self, batch, batch_idx):
input_ids, attention_mask, labels = batch["input_ids"], batch["attention_mask"], batch["label"]
logits = self(input_ids, attention_mask)
loss = self.loss_fn(logits, labels)
preds = torch.argmax(logits, dim=1)
acc = accuracy_score(labels.cpu(), preds.cpu())
f1 = f1_score(labels.cpu(), preds.cpu(), average='weighted')
self.log("val_loss", loss, prog_bar=True)
self.log("val_acc", acc, prog_bar=True)
self.log("val_f1", f1, prog_bar=True, sync_dist=True)
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=self.hparams.lr)
if __name__ == "__main__":
df = pd.read_csv("data_cleaned2.csv")
class_counts = df["label"].value_counts().sort_index().tolist()
class_weights = 1.0 / np.array(class_counts)
class_weights = class_weights / class_weights.sum()
train_df = df.sample(frac=0.8, random_state=42)
val_df = df.drop(train_df.index)
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
def tokenize(batch):
return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=64)
train_dataset = Dataset.from_pandas(train_df).map(tokenize, batched=True)
val_dataset = Dataset.from_pandas(val_df).map(tokenize, batched=True)
train_dataset.set_format("torch", columns=["input_ids", "attention_mask", "label"])
val_dataset.set_format("torch", columns=["input_ids", "attention_mask", "label"])
train_loader = DataLoader(train_dataset, batch_size=16, num_workers=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16, num_workers=8)
checkpoint_callback = ModelCheckpoint(
dirpath="checkpoints/",
filename="roberta-priority-{epoch:02d}-{val_f1:.2f}",
save_top_k=3,
monitor="val_f1",
mode="max",
save_weights_only=True,
every_n_epochs=1
)
early_stopping = EarlyStopping(
monitor="val_f1",
patience=2,
mode="max",
verbose=True,
)
trainer_kwargs = dict(
accelerator="gpu",
devices=2,
strategy=DDPStrategy(find_unused_parameters=True),
max_epochs=20,
precision=16,
log_every_n_steps=10,
callbacks=[checkpoint_callback, early_stopping],
)
trainer = pl.Trainer(**trainer_kwargs)
model = RoBERTaClassifier(class_weights=class_weights)
trainer.fit(model, train_loader, val_loader)
|