Upload 4 files

eurovoc.py

@@ -0,0 +1,212 @@
+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)
+ 

handler.py

@@ -0,0 +1,75 @@
+from typing import Dict, List, Any
+import numpy as np
+import pickle
+
+from sklearn.preprocessing import MultiLabelBinarizer
+from transformers import AutoTokenizer
+import torch
+
+from eurovoc import EurovocTagger
+
+BERT_MODEL_NAME = "nlpaueb/legal-bert-base-uncased"
+MAX_LEN = 512
+TEXT_MAX_LEN = MAX_LEN * 50
+tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL_NAME)
+
+
+class EndpointHandler:
+ mlb = MultiLabelBinarizer()
+
+ def __init__(self, path=""):
+ self.mlb = pickle.load(open(f"{path}/mlb.pickle", "rb"))
+
+ self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+ self.model = EurovocTagger.from_pretrained(path,
+ bert_model_name=BERT_MODEL_NAME,
+ n_classes=len(self.mlb.classes_),
+ map_location=self.device)
+ self.model.eval()
+ self.model.freeze()
+
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+ """
+ data args:
+ inputs (:obj: `str` | `PIL.Image` | `np.array`)
+ kwargs
+ Return:
+ A :obj:`list` | `dict`: will be serialized and returned
+ """
+
+ text = data.pop("inputs", data)
+ topk = data.pop("topk", 5)
+ threshold = data.pop("threshold", 0.16)
+ debug = data.pop("debug", False)
+ prediction = self.get_prediction(text)
+ results = [{"label": label, "score": float(score)} for label, score in
+ zip(self.mlb.classes_, prediction[0].tolist())]
+ results = sorted(results, key=lambda x: x["score"], reverse=True)
+ results = [r for r in results if r["score"] > threshold]
+ results = results[:topk]
+ if debug:
+ return {"results": results, "values": prediction, "input": text}
+ else:
+ return {"results": results}
+
+ def get_prediction(self, text):
+ # split text into chunks of MAX_LEN and get average prediction for each chunk
+ chunks = [text[i:i + MAX_LEN] for i in range(0, min(len(text), TEXT_MAX_LEN), MAX_LEN)]
+ predictions = [self._get_prediction(chunk) for chunk in chunks]
+ predictions = np.array(predictions).mean(axis=0)
+ return predictions
+
+ def _get_prediction(self, text):
+ item = tokenizer.encode_plus(
+ text,
+ add_special_tokens=True,
+ max_length=MAX_LEN,
+ return_token_type_ids=False,
+ padding="max_length",
+ truncation=True,
+ return_attention_mask=True,
+ return_tensors='pt')
+ _, prediction = self.model(item["input_ids"], item["attention_mask"])
+ prediction = prediction.cpu().detach().numpy()
+ return prediction
+

mlb.pickle

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:35015ecbd09a8524d555feb303f81788fc8be9dd28ae2eae9f4e05f7417b1d71
+size 122082

requirements.txt

@@ -0,0 +1,8 @@
+datasets==2.13.1
+ipykernel==6.24.0
+lightning==2.0.5
+pip-chill==1.0.3
+scikit-learn==1.3.0
+scikit-multilearn==0.2.0
+transformers==4.30.2
+