eurovoc_fr / eurovoc.py
scampion's picture
Upload 4 files
c480b17
raw
history blame
6.55 kB
import torch
from torch.utils.data import Dataset, DataLoader
import numpy as np
import pytorch_lightning as pl
import torch.nn as nn
from transformers import BertTokenizerFast as BertTokenizer, AdamW, get_linear_schedule_with_warmup, AutoTokenizer, AutoModel
from huggingface_hub import PyTorchModelHubMixin
class EurovocDataset(Dataset):
def __init__(
self,
text: np.array,
labels: np.array,
tokenizer: BertTokenizer,
max_token_len: int = 128
):
self.tokenizer = tokenizer
self.text = text
self.labels = labels
self.max_token_len = max_token_len
def __len__(self):
return len(self.labels)
def __getitem__(self, index: int):
text = self.text[index][0]
labels = self.labels[index]
encoding = self.tokenizer.encode_plus(
text,
add_special_tokens=True,
max_length=self.max_token_len,
return_token_type_ids=False,
padding="max_length",
truncation=True,
return_attention_mask=True,
return_tensors='pt',
)
return dict(
text=text,
input_ids=encoding["input_ids"].flatten(),
attention_mask=encoding["attention_mask"].flatten(),
labels=torch.FloatTensor(labels)
)
class EuroVocLongTextDataset(Dataset):
def __splitter__(text, max_lenght):
l = text.split()
for i in range(0, len(l), max_lenght):
yield l[i:i + max_lenght]
def __init__(
self,
text: np.array,
labels: np.array,
tokenizer: BertTokenizer,
max_token_len: int = 128
):
self.tokenizer = tokenizer
self.text = text
self.labels = labels
self.max_token_len = max_token_len
self.chunks_and_labels = [(c, l) for t, l in zip(self.text, self.labels) for c in self.__splitter__(t)]
self.encoding = self.tokenizer.batch_encode_plus(
[c for c, _ in self.chunks_and_labels],
add_special_tokens=True,
max_length=self.max_token_len,
return_token_type_ids=False,
padding="max_length",
truncation=True,
return_attention_mask=True,
return_tensors='pt',
)
def __len__(self):
return len(self.chunks_and_labels)
def __getitem__(self, index: int):
text, labels = self.chunks_and_labels[index]
return dict(
text=text,
input_ids=self.encoding[index]["input_ids"].flatten(),
attention_mask=self.encoding[index]["attention_mask"].flatten(),
labels=torch.FloatTensor(labels)
)
class EurovocDataModule(pl.LightningDataModule):
def __init__(self, bert_model_name, x_tr, y_tr, x_test, y_test, batch_size=8, max_token_len=512):
super().__init__()
self.batch_size = batch_size
self.x_tr = x_tr
self.y_tr = y_tr
self.x_test = x_test
self.y_test = y_test
self.tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
self.max_token_len = max_token_len
def setup(self, stage=None):
self.train_dataset = EurovocDataset(
self.x_tr,
self.y_tr,
self.tokenizer,
self.max_token_len
)
self.test_dataset = EurovocDataset(
self.x_test,
self.y_test,
self.tokenizer,
self.max_token_len
)
def train_dataloader(self):
return DataLoader(
self.train_dataset,
batch_size=self.batch_size,
shuffle=True,
num_workers=2
)
def val_dataloader(self):
return DataLoader(
self.test_dataset,
batch_size=self.batch_size,
num_workers=2
)
def test_dataloader(self):
return DataLoader(
self.test_dataset,
batch_size=self.batch_size,
num_workers=2
)
class EurovocTagger(pl.LightningModule, PyTorchModelHubMixin):
def __init__(self, bert_model_name, n_classes, lr=2e-5, eps=1e-8):
super().__init__()
self.bert = AutoModel.from_pretrained(bert_model_name)
self.dropout = nn.Dropout(p=0.2)
self.classifier1 = nn.Linear(self.bert.config.hidden_size, n_classes)
self.criterion = nn.BCELoss()
self.lr = lr
self.eps = eps
def forward(self, input_ids, attention_mask, labels=None):
output = self.bert(input_ids, attention_mask=attention_mask)
output = self.dropout(output.pooler_output)
output = self.classifier1(output)
output = torch.sigmoid(output)
loss = 0
if labels is not None:
loss = self.criterion(output, labels)
return loss, output
def training_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
self.log("train_loss", loss, prog_bar=True, logger=True)
return {"loss": loss, "predictions": outputs, "labels": labels}
def validation_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
self.log("val_loss", loss, prog_bar=True, logger=True)
return loss
def test_step(self, batch, batch_idx):
input_ids = batch["input_ids"]
attention_mask = batch["attention_mask"]
labels = batch["labels"]
loss, outputs = self(input_ids, attention_mask, labels)
self.log("test_loss", loss, prog_bar=True, logger=True)
return loss
def on_train_epoch_end(self, *args, **kwargs):
return
#labels = []
#predictions = []
#for output in args['outputs']:
# for out_labels in output["labels"].detach().cpu():
# labels.append(out_labels)
# for out_predictions in output["predictions"].detach().cpu():
# predictions.append(out_predictions)
#labels = torch.stack(labels).int()
#predictions = torch.stack(predictions)
#for i, name in enumerate(mlb.classes_):
# class_roc_auc = auroc(predictions[:, i], labels[:, i])
# self.logger.experiment.add_scalar(f"{name}_roc_auc/Train", class_roc_auc, self.current_epoch)
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=self.lr, eps=self.eps)