train bert on NER

Author: yan-vei

.gitattributes

@@ -0,0 +1 @@
+*.arrow filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -1,2 +1,3 @@
 .env
-.idea
+.idea
+__pycache__/

datasets/kaznerd-test.hf/data-00000-of-00001.arrow

@@ -0,0 +1,42 @@
+import torch
+import torch.nn as nn
+from datasets import load_from_disk
+from torch.utils.data import DataLoader
+from utils.tokenizer import NERTokenizer
+from utils.train import train_ner
+from models.base_model import BertNerd
+
+config = {
+    'PADDING_TOKEN': -100,
+    'LEARNING_RATE': 0.001,
+    'NUM_EPOCHS': 10,
+    'BATCH_SIZE': 16,
+    'RANDOM_SEED': 42,
+    'CHUNK_SIZE': 100,
+    'HIDDEN_SIZE': 768
+}
+
+# Load datasets
+kaznerd_train = load_from_disk('datasets/kaznerd-train.hf')
+kaznerd_test = load_from_disk('datasets/kaznerd-test.hf')
+
+kz_labels_list = kaznerd_train.features["ner_tags"].feature.names
+config['NUM_CLASSES'] = len(kz_labels_list)
+config['DEVICE'] = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Initialize tokenizer
+tokenizer = NERTokenizer("bert-base-uncased")
+
+# Tokenize and create dataloaders for Kazakh NER dataset
+kz_tokenized_train = kaznerd_train.map(lambda e: tokenizer.tokenize_and_align_labels(e, tags='ner_tags'), batched=True, batch_size=config['BATCH_SIZE'])
+kz_tokenized_train.set_format(type='torch', columns=['input_ids', 'attention_mask', 'labels'])
+
+kz_train_dataloader = DataLoader(kz_tokenized_train, batch_size=config['BATCH_SIZE'])
+
+# Define model, loss function, optimizer
+kaznerd_model = BertNerd(config)
+loss_func = nn.CrossEntropyLoss(ignore_index=config['PADDING_TOKEN'])
+optimizer = torch.optim.Adam(kaznerd_model.get_params(), lr=config['LEARNING_RATE'])
+
+train_ner(model=kaznerd_model, optimizer=optimizer, loss_func=loss_func, train_dataloader=kz_train_dataloader,
+          config=config)

datasets/kaznerd-train.hf/data-00000-of-00001.arrow

@@ -1,21 +1,23 @@
 import torch
-from transformers import AutoModelForMaskedLM
+from transformers import BertModel
 
 class BertNerd(torch.nn.Module):
     """
         MBert-based model for performing NER tasks w/ and w/o
         soft prompts on Kazakh and Turkish languages.
     """
 
-    def __init__(self, config, device, freeze=True):
+    def __init__(self, config, freeze=True):
         super(BertNerd, self).__init__()
-        self.mbert = AutoModelForMaskedLM("google-bert/bert-base-multilingual-cased")
-        self.linear = torch.nn.Linear(config.hidden_size, config.num_classes)
-        self.device = device
+        self.device = config['DEVICE']
+        self.mbert = BertModel.from_pretrained("google-bert/bert-base-multilingual-uncased").to(self.device)
+        self.linear = torch.nn.Linear(config['HIDDEN_SIZE'], config['NUM_CLASSES'])
 
         if freeze:
             self.freeze_params()
 
+        print("\tModel initilized.")
+
     def forward(self, input_seq, attention_mask):
         """
         Define the model's forward pass.
@@ -24,37 +26,32 @@ def forward(self, input_seq, attention_mask):
         :param attention_mask: attention mask
         :return: predicted logits
         """
-
         input_seq = self.mbert(input_seq, attention_mask).last_hidden_state.to(self.device)
         logits = self.linear(input_seq)
 
         return logits
 
-    def get_loss(self, loss_fn, logits, labels, ignore_index=None):
+    def freeze_params(self):
         """
-        Get loss for the forward pass of the current batch.
+        Only train the soft prompts, don't train any model parameters.
 
-        :param loss_fn: e.g. nn.CrossEntropyLoss
-        :param logits: predicted labels
-        :param labels: actual labels from the dataset
-        :param ignore_index: padding index to ignore
-        :return: loss per batch
+        :return: void
         """
 
-        loss_func = loss_fn(ignore_index=ignore_index)
-
-        # ToDo
-        # Logits/labels should probably be flattened, so we get the right dimension
-
-        return loss_func(logits, labels).detach().item()
-
+        for param in self.mbert.parameters():
+            param.requires_grad = False
 
-    def freeze_params(self):
+    def get_params(self):
         """
-        Only train the soft prompts, don't train any model parameters.
+        Return tunable parameters of the model.
 
-        :return: void
+        :return: list of tunable params
         """
 
-        for param in self.bert.parameters():
-            param.requires_grad = False
+        params = []
+
+        for param in self.parameters():
+            if param.requires_grad:
+                params.append(param)
+
+        return params

datasets/kaznerd-val.hf/data-00000-of-00001.arrow

@@ -1,18 +1,22 @@
 import torch
-from metrics import get_accuracy
+from .metrics import get_accuracy
 
 
-def train_ner(model, train_dataloader, optimizer, config):
+def train_ner(model, train_dataloader, loss_func, optimizer, config):
     """
         Define the training loop for NER.
     :param model: corresponding model class
-    :param train_loader: train data
+    :param train_dataloader: train data
+    :param loss_func: loss function
+    :param optimizer: optimizer
     :param config: config file with hyperparameters
     :return: model, metrics
     """
+    print("\tTraining started.")
+
     accuracies = []
 
-    for epoch in range(config.epochs):
+    for epoch in range(config['NUM_EPOCHS']):
         loss_per_epoch = 0
         correct = 0
         total = 0
@@ -21,14 +25,15 @@ def train_ner(model, train_dataloader, optimizer, config):
         model.train()
 
         for batch in train_dataloader:
-            inputs, attention_mask, labels = batch["input_ids"].to(config.device), batch["attention_mask"].to(config.device), batch["labels"].to(config.device)
+            inputs, attention_mask, labels = (batch["input_ids"].to(config['DEVICE']), batch["attention_mask"].to(config['DEVICE']),
+                                              batch["labels"].to(config['DEVICE']))
 
             # Make prediction
             logits = model(inputs, attention_mask)
 
             # Calculate loss
-            batch_loss = model.get_loss(logits, labels)
-            loss_per_epoch += batch_loss
+            batch_loss = loss_func(logits.flatten(end_dim=1), labels.flatten(end_dim=1))
+            loss_per_epoch += batch_loss.detach().item()
 
             # Get ids corresponding to the most probably NER tags
             tag_ids = torch.max(logits, dim=2).indices