SinaTools 0.1.41__py2.py3-none-any.whl → 1.0.1__py2.py3-none-any.whl

{SinaTools-0.1.41.dist-info → SinaTools-1.0.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: SinaTools
-Version: 0.1.41
+Version: 1.0.1
 Summary: Open-source Python toolkit for Arabic Natural Understanding, allowing people to integrate it in their system workflow.
 Home-page: https://github.com/SinaLab/sinatools
 License: MIT license

SinaTools-1.0.1.dist-info/RECORD

@@ -0,0 +1,73 @@
+SinaTools-1.0.1.data/data/sinatools/environment.yml,sha256=i0UFZc-vwU9ZwnI8hBdz7vi-x22vG-HR8ojWBUAOkno,5422
+sinatools/VERSION,sha256=1R5uyUBYVUqEVYpbQC7m71_fVFXjXJAv7aYc2odSlDo,5
+sinatools/__init__.py,sha256=bEosTU1o-FSpyytS6iVP_82BXHF2yHnzpJxPLYRbeII,135
+sinatools/environment.yml,sha256=i0UFZc-vwU9ZwnI8hBdz7vi-x22vG-HR8ojWBUAOkno,5422
+sinatools/install_env.py,sha256=EODeeE0ZzfM_rz33_JSIruX03Nc4ghyVOM5BHVhsZaQ,404
+sinatools/sinatools.py,sha256=vR5AaF0iel21LvsdcqwheoBz0SIj9K9I_Ub8M8oA98Y,20
+sinatools/CLI/DataDownload/download_files.py,sha256=EezvbukR3pZ8s6mGZnzTcjsbo3CBDlC0g6KhJWlYp1w,2686
+sinatools/CLI/morphology/ALMA_multi_word.py,sha256=rmpa72twwIJHme_kpQ1lu3_7y_Jorj70QTvOnQMJRuI,1274
+sinatools/CLI/morphology/morph_analyzer.py,sha256=HPamEKos_JRYCJv_2q6c12N--da58_JXTno9haww5Ao,3497
+sinatools/CLI/ner/corpus_entity_extractor.py,sha256=DdvigsDQzko5nJBjzUXlIDqoBMBTVzktjSo7JfEXTIA,4778
+sinatools/CLI/ner/entity_extractor.py,sha256=G9j-t0WKm2CRORhqARJM-pI-KArQ2IXIvnBK_NHxlHs,2885
+sinatools/CLI/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sinatools/CLI/utils/arStrip.py,sha256=NLyp8vOu2xv80tL9jiKRvyptmbkRZVg-wcAr-9YyvNY,3264
+sinatools/CLI/utils/corpus_tokenizer.py,sha256=nH0T4h6urr_0Qy6-wN3PquOtnwybj0REde5Ts_OE4U8,1650
+sinatools/CLI/utils/implication.py,sha256=AojpkCwUQJiQjxhyEUWKRHmBnIt1tVqr485cAF7Thq0,2857
+sinatools/CLI/utils/jaccard.py,sha256=w56N_cNEFJ0A7WtunmY_xtms4srFagKBzrW_0YhH2DE,4216
+sinatools/CLI/utils/remove_latin.py,sha256=NOaTm2RHxt5IQrV98ySTmD8rTXTmcqSmfbPAwTyaXqU,848
+sinatools/CLI/utils/remove_punctuation.py,sha256=vJAZlEn7WGftZAFVFYnddkRrxdJ_rMmKB9vFZkY-jN4,1097
+sinatools/CLI/utils/sentence_tokenizer.py,sha256=Wli8eiDbWSd_Z8UKpu_JkaS8jImowa1vnRL0oYCSfqw,2823
+sinatools/CLI/utils/text_dublication_detector.py,sha256=dW70O5O20GxeUDDF6zVYn52wWLmJF-HBZgvqIeVL2rQ,1661
+sinatools/CLI/utils/text_transliteration.py,sha256=vz-3kxWf8pNYVCqNAtBAiA6u_efrS5NtWT-ofN1NX6I,2014
+sinatools/DataDownload/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sinatools/DataDownload/downloader.py,sha256=VdUNgSqMKz1J-DuQD_eS1U2KWqEpy94WlSJ0pPODLig,7833
+sinatools/arabert/__init__.py,sha256=ely2PttjgSv7vKdzskuD1rtK_l_UOpmxJSz8isrveD0,16
+sinatools/arabert/preprocess.py,sha256=qI0FsuMTOzdRlYGCtLrjpXgikNElUZPv9bnjaKDZKJ4,33024
+sinatools/morphology/ALMA_multi_word.py,sha256=hj_-8ojrYYHnfCGk8WKtJdUR8mauzQdma4WUm-okDps,1346
+sinatools/morphology/__init__.py,sha256=I4wVBh8BhyNl-CySVdiI_nUSn6gj1j-gmLKP300RpE0,1216
+sinatools/morphology/morph_analyzer.py,sha256=JOH2UWKNQWo5UzpWNzP9R1D3B3qLSogIiMp8n0N_56o,7177
+sinatools/ner/__init__.py,sha256=59kLMX6UQhF6JpE10RhaDYC3a2_jiWOIVPuejsoflFE,1050
+sinatools/ner/data_format.py,sha256=VmFshZbEPOsWxsb4tgSkwvbM1k7yCce4kmtPkCiWgwM,4513
+sinatools/ner/datasets.py,sha256=mG1iwqSm3lXCFHLqE-b4wNi176cpuzNBz8tKaBU6z6M,5059
+sinatools/ner/entity_extractor.py,sha256=O2epRwRFUUcQs3SnFIYHVBI4zVhr8hRcj0XJYeby4ts,3588
+sinatools/ner/helpers.py,sha256=sX6ezVbuVQxk_xJqZwhUzJVFVuVmFGmei_kd6r3sPHE,3652
+sinatools/ner/metrics.py,sha256=Irz6SsIvpOzGIA2lWxrEV86xnTnm0TzKm9SUVT4SXUU,2734
+sinatools/ner/transforms.py,sha256=vti3mDdi-IRP8i0aTQ37QqpPlP9hdMmJ6_bAMa0uL-s,4871
+sinatools/ner/data/__init__.py,sha256=W0C1ge_XxTfmdEGz0hkclz57aLI5VFS5t6BjByCfkFk,57
+sinatools/ner/data/datasets.py,sha256=_uUlvBAhnTtPwKLj0wIbmB04VCBidfwffxKorLGHq_g,5134
+sinatools/ner/data/transforms.py,sha256=URMz1dHzkHjgUGAkDOenCWvQThO1ha8XeQVjoLL9RXM,4874
+sinatools/ner/nn/BaseModel.py,sha256=3GmujQasTZZunOBuFXpY2p1W8W256iI_Uu4hxhOY2Z0,608
+sinatools/ner/nn/BertNestedTagger.py,sha256=_fwAn1kiKmXe6m5y16Ipty3kvXIEFEmiUq74Ad1818U,1219
+sinatools/ner/nn/BertSeqTagger.py,sha256=dFcBBiMw2QCWsyy7aQDe_PS3aRuNn4DOxKIHgTblFvc,504
+sinatools/ner/nn/__init__.py,sha256=UgQD_XLNzQGBNSYc_Bw1aRJZjq4PJsnMT1iZwnJemqE,170
+sinatools/ner/trainers/BaseTrainer.py,sha256=Uar8HxtgBXCVhKa85sEN622d9P7JiFBcWfs46uRG4aA,4068
+sinatools/ner/trainers/BertNestedTrainer.py,sha256=Pb4O2WeBmTvV3hHMT6DXjxrTzgtuh3OrKQZnogYy8RQ,8429
+sinatools/ner/trainers/BertTrainer.py,sha256=B_uVtUwfv_eFwMMPsKQvZgW_ZNLy6XEsX5ePR0s8d-k,6433
+sinatools/ner/trainers/__init__.py,sha256=UDok8pDDpYOpwRBBKVLKaOgSUlmqqb-zHZI1p0xPxzI,188
+sinatools/relations/__init__.py,sha256=cYjsP2mlTYvAwVIEFtgA6i9gLUSkGVOuDggMs7TvG5k,272
+sinatools/relations/relation_extractor.py,sha256=UuDlaaR0ch9BFv4sBF1tr7P-P9xq8oRZF41tAze6_ok,9751
+sinatools/semantic_relatedness/__init__.py,sha256=S0xrmqtl72L02N56nbNMudPoebnYQgsaIyyX-587DsU,830
+sinatools/semantic_relatedness/compute_relatedness.py,sha256=_9HFPs3nQBLklHFfkc9o3gEjEI6Bd34Ha4E1Kvv1RIg,2256
+sinatools/synonyms/__init__.py,sha256=yMuphNZrm5XLOR2T0weOHcUysJm-JKHUmVLoLQO8390,548
+sinatools/synonyms/synonyms_generator.py,sha256=jRd0D3_kn-jYBaZzqY-7oOy0SFjSJ-mjM7JhsySzX58,9037
+sinatools/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+sinatools/utils/charsets.py,sha256=rs82oZJqRqosZdTKXfFAJfJ5t4PxjMM_oAPsiWSWuwU,2817
+sinatools/utils/parser.py,sha256=qvHdln5R5CAv_0UOJWe0mcp8JCsGqgazoeIIkoALH88,6259
+sinatools/utils/readfile.py,sha256=xE4LEaCqXJIk9v37QUSSmWb-aY3UnCFUNb7uVdx3cpM,133
+sinatools/utils/similarity.py,sha256=HAK6OmyVnfjPm0GWL3z9s4ZoUwpZHVKxt3CeSMfqLIQ,11990
+sinatools/utils/text_dublication_detector.py,sha256=FeSkbfWGMQluz23H4CBHXION-walZPgjueX6AL8u_Q0,5660
+sinatools/utils/text_transliteration.py,sha256=F3smhr2AEJtySE6wGQsiXXOslTvSDzLivTYu0btgc10,8769
+sinatools/utils/tokenizer.py,sha256=nyk6lh5-p38wrU62hvh4wg7ni9ammkdqqIgcjbbBxxo,6965
+sinatools/utils/tokenizers_words.py,sha256=efNfOil9qDNVJ9yynk_8sqf65PsL-xtsHG7y2SZCkjQ,656
+sinatools/utils/word_compare.py,sha256=rS2Z74sf7R-7MTXyrFj5miRi2TnSG9OdTDp_qQYuo2Y,28200
+sinatools/wsd/__init__.py,sha256=mwmCUurOV42rsNRpIUP3luG0oEzeTfEx3oeDl93Oif8,306
+sinatools/wsd/disambiguator.py,sha256=h-3idc5rPPbMDSE_QVJAsEVkDHwzYY3L2SEPNXIdOcc,20104
+sinatools/wsd/settings.py,sha256=6XflVTFKD8SVySX9Wj7zYQtV26WDTcQ2-uW8-gDNHKE,747
+sinatools/wsd/wsd.py,sha256=gHIBUFXegoY1z3rRnIlK6TduhYq2BTa_dHakOjOlT4k,4434
+SinaTools-1.0.1.dist-info/AUTHORS.rst,sha256=aTWeWlIdfLi56iLJfIUAwIrmqDcgxXKLji75_Fjzjyg,174
+SinaTools-1.0.1.dist-info/LICENSE,sha256=uwsKYG4TayHXNANWdpfMN2lVW4dimxQjA_7vuCVhD70,1088
+SinaTools-1.0.1.dist-info/METADATA,sha256=8EnFO3dSqtJ8JJ4r_-ji5tX_h04_vNTnPvfubqceaQ4,3409
+SinaTools-1.0.1.dist-info/WHEEL,sha256=9Hm2OB-j1QcCUq9Jguht7ayGIIZBRTdOXD1qg9cCgPM,109
+SinaTools-1.0.1.dist-info/entry_points.txt,sha256=_CsRKM_tSCWV5hefBNUsWf9_6DrJnzFlxeAo1wm5XqY,1302
+SinaTools-1.0.1.dist-info/top_level.txt,sha256=8tNdPTeJKw3TQCaua8IJIx6N6WpgZZmVekf1OdBNJpE,10
+SinaTools-1.0.1.dist-info/RECORD,,

sinatools/VERSION

@@ -1 +1 @@
-0.1.41
+1.0.1

sinatools/ner/trainers/BertNestedTrainer.py

@@ -1,203 +1,203 @@
-import os
-import logging
-import torch
-import numpy as np
-from sinatools.ner.trainers import BaseTrainer
-from sinatools.ner.metrics import compute_nested_metrics
-
-logger = logging.getLogger(__name__)
-
-
-class BertNestedTrainer(BaseTrainer):
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
-
-    def train(self):
-        best_val_loss, test_loss = np.inf, np.inf
-        num_train_batch = len(self.train_dataloader)
-        num_labels = [len(v) for v in self.train_dataloader.dataset.vocab.tags[1:]]
-        patience = self.patience
-
-        for epoch_index in range(self.max_epochs):
-            self.current_epoch = epoch_index
-            train_loss = 0
-
-            for batch_index, (subwords, gold_tags, tokens, valid_len, logits) in enumerate(self.tag(
-                self.train_dataloader, is_train=True
-            ), 1):
-                self.current_timestep += 1
-
-                # Compute loses for each output
-                # logits = B x T x L x C
-                losses = [self.loss(logits[:, :, i, 0:l].view(-1, logits[:, :, i, 0:l].shape[-1]),
-                                    torch.reshape(gold_tags[:, i, :], (-1,)).long())
-                          for i, l in enumerate(num_labels)]
-
-                torch.autograd.backward(losses)
-
-                # Avoid exploding gradient by doing gradient clipping
-                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
-
-                self.optimizer.step()
-                self.scheduler.step()
-                batch_loss = sum(l.item() for l in losses)
-                train_loss += batch_loss
-
-                if self.current_timestep % self.log_interval == 0:
-                    logger.info(
-                        "Epoch %d | Batch %d/%d | Timestep %d | LR %.10f | Loss %f",
-                        epoch_index,
-                        batch_index,
-                        num_train_batch,
-                        self.current_timestep,
-                        self.optimizer.param_groups[0]['lr'],
-                        batch_loss
-                    )
-
-            train_loss /= num_train_batch
-
-            logger.info("** Evaluating on validation dataset **")
-            val_preds, segments, valid_len, val_loss = self.eval(self.val_dataloader)
-            val_metrics = compute_nested_metrics(segments, self.val_dataloader.dataset.transform.vocab.tags[1:])
-
-            epoch_summary_loss = {
-                "train_loss": train_loss,
-                "val_loss": val_loss
-            }
-            epoch_summary_metrics = {
-                "val_micro_f1": val_metrics.micro_f1,
-                "val_precision": val_metrics.precision,
-                "val_recall": val_metrics.recall
-            }
-
-            logger.info(
-                "Epoch %d | Timestep %d | Train Loss %f | Val Loss %f | F1 %f",
-                epoch_index,
-                self.current_timestep,
-                train_loss,
-                val_loss,
-                val_metrics.micro_f1
-            )
-
-            if val_loss < best_val_loss:
-                patience = self.patience
-                best_val_loss = val_loss
-                logger.info("** Validation improved, evaluating test data **")
-                test_preds, segments, valid_len, test_loss = self.eval(self.test_dataloader)
-                self.segments_to_file(segments, os.path.join(self.output_path, "predictions.txt"))
-                test_metrics = compute_nested_metrics(segments, self.test_dataloader.dataset.transform.vocab.tags[1:])
-
-                epoch_summary_loss["test_loss"] = test_loss
-                epoch_summary_metrics["test_micro_f1"] = test_metrics.micro_f1
-                epoch_summary_metrics["test_precision"] = test_metrics.precision
-                epoch_summary_metrics["test_recall"] = test_metrics.recall
-
-                logger.info(
-                    f"Epoch %d | Timestep %d | Test Loss %f | F1 %f",
-                    epoch_index,
-                    self.current_timestep,
-                    test_loss,
-                    test_metrics.micro_f1
-                )
-
-                self.save()
-            else:
-                patience -= 1
-
-            # No improvements, terminating early
-            if patience == 0:
-                logger.info("Early termination triggered")
-                break
-
-            self.summary_writer.add_scalars("Loss", epoch_summary_loss, global_step=self.current_timestep)
-            self.summary_writer.add_scalars("Metrics", epoch_summary_metrics, global_step=self.current_timestep)
-
-    def tag(self, dataloader, is_train=True):
-        """
-        Given a dataloader containing segments, predict the tags
-        :param dataloader: torch.utils.data.DataLoader
-        :param is_train: boolean - True for training model, False for evaluation
-        :return: Iterator
-                    subwords (B x T x NUM_LABELS)- torch.Tensor - BERT subword ID
-                    gold_tags (B x T x NUM_LABELS) - torch.Tensor - ground truth tags IDs
-                    tokens - List[arabiner.data.dataset.Token] - list of tokens
-                    valid_len (B x 1) - int - valiud length of each sequence
-                    logits (B x T x NUM_LABELS) - logits for each token and each tag
-        """
-        for subwords, gold_tags, tokens, mask, valid_len in dataloader:
-            self.model.train(is_train)
-
-            if torch.cuda.is_available():
-                subwords = subwords.cuda()
-                gold_tags = gold_tags.cuda()
-
-            if is_train:
-                self.optimizer.zero_grad()
-                logits = self.model(subwords)
-            else:
-                with torch.no_grad():
-                    logits = self.model(subwords)
-
-            yield subwords, gold_tags, tokens, valid_len, logits
-
-    def eval(self, dataloader):
-        golds, preds, segments, valid_lens = list(), list(), list(), list()
-        num_labels = [len(v) for v in dataloader.dataset.vocab.tags[1:]]
-        loss = 0
-
-        for _, gold_tags, tokens, valid_len, logits in self.tag(
-            dataloader, is_train=False
-        ):
-            losses = [self.loss(logits[:, :, i, 0:l].view(-1, logits[:, :, i, 0:l].shape[-1]),
-                                torch.reshape(gold_tags[:, i, :], (-1,)).long())
-                      for i, l in enumerate(num_labels)]
-            loss += sum(losses)
-            preds += torch.argmax(logits, dim=3)
-            segments += tokens
-            valid_lens += list(valid_len)
-
-        loss /= len(dataloader)
-
-        # Update segments, attach predicted tags to each token
-        segments = self.to_segments(segments, preds, valid_lens, dataloader.dataset.vocab)
-
-        return preds, segments, valid_lens, loss
-
-    def infer(self, dataloader):
-        golds, preds, segments, valid_lens = list(), list(), list(), list()
-
-        for _, gold_tags, tokens, valid_len, logits in self.tag(
-            dataloader, is_train=False
-        ):
-            preds += torch.argmax(logits, dim=3)
-            segments += tokens
-            valid_lens += list(valid_len)
-
-        segments = self.to_segments(segments, preds, valid_lens, dataloader.dataset.vocab)
-        return segments
-
-    def to_segments(self, segments, preds, valid_lens, vocab):
-        if vocab is None:
-            vocab = self.vocab
-
-        tagged_segments = list()
-        tokens_stoi = vocab.tokens.get_stoi()
-        unk_id = tokens_stoi["UNK"]
-
-        for segment, pred, valid_len in zip(segments, preds, valid_lens):
-            # First, the token at 0th index [CLS] and token at nth index [SEP]
-            # Combine the tokens with their corresponding predictions
-            segment_pred = zip(segment[1:valid_len-1], pred[1:valid_len-1])
-
-            # Ignore the sub-tokens/subwords, which are identified with text being UNK
-            segment_pred = list(filter(lambda t: tokens_stoi[t[0].text] != unk_id, segment_pred))
-
-            # Attach the predicted tags to each token
-            list(map(lambda t: setattr(t[0], 'pred_tag', [{"tag": vocab.get_itos()[tag_id]}
-                                                     for tag_id, vocab in zip(t[1].int().tolist(), vocab.tags[1:])]), segment_pred))
-
-            # We are only interested in the tagged tokens, we do no longer need raw model predictions
-            tagged_segment = [t for t, _ in segment_pred]
-            tagged_segments.append(tagged_segment)
-
-        return tagged_segments
+import os
+import logging
+import torch
+import numpy as np
+from sinatools.ner.trainers import BaseTrainer
+from sinatools.ner.metrics import compute_nested_metrics
+
+logger = logging.getLogger(__name__)
+
+
+class BertNestedTrainer(BaseTrainer):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+
+    def train(self):
+        best_val_loss, test_loss = np.inf, np.inf
+        num_train_batch = len(self.train_dataloader)
+        num_labels = [len(v) for v in self.train_dataloader.dataset.vocab.tags[1:]]
+        patience = self.patience
+
+        for epoch_index in range(self.max_epochs):
+            self.current_epoch = epoch_index
+            train_loss = 0
+
+            for batch_index, (subwords, gold_tags, tokens, valid_len, logits) in enumerate(self.tag(
+                self.train_dataloader, is_train=True
+            ), 1):
+                self.current_timestep += 1
+
+                # Compute loses for each output
+                # logits = B x T x L x C
+                losses = [self.loss(logits[:, :, i, 0:l].view(-1, logits[:, :, i, 0:l].shape[-1]),
+                                    torch.reshape(gold_tags[:, i, :], (-1,)).long())
+                          for i, l in enumerate(num_labels)]
+
+                torch.autograd.backward(losses)
+
+                # Avoid exploding gradient by doing gradient clipping
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
+
+                self.optimizer.step()
+                self.scheduler.step()
+                batch_loss = sum(l.item() for l in losses)
+                train_loss += batch_loss
+
+                if self.current_timestep % self.log_interval == 0:
+                    logger.info(
+                        "Epoch %d | Batch %d/%d | Timestep %d | LR %.10f | Loss %f",
+                        epoch_index,
+                        batch_index,
+                        num_train_batch,
+                        self.current_timestep,
+                        self.optimizer.param_groups[0]['lr'],
+                        batch_loss
+                    )
+
+            train_loss /= num_train_batch
+
+            logger.info("** Evaluating on validation dataset **")
+            val_preds, segments, valid_len, val_loss = self.eval(self.val_dataloader)
+            val_metrics = compute_nested_metrics(segments, self.val_dataloader.dataset.transform.vocab.tags[1:])
+
+            epoch_summary_loss = {
+                "train_loss": train_loss,
+                "val_loss": val_loss
+            }
+            epoch_summary_metrics = {
+                "val_micro_f1": val_metrics.micro_f1,
+                "val_precision": val_metrics.precision,
+                "val_recall": val_metrics.recall
+            }
+
+            logger.info(
+                "Epoch %d | Timestep %d | Train Loss %f | Val Loss %f | F1 %f",
+                epoch_index,
+                self.current_timestep,
+                train_loss,
+                val_loss,
+                val_metrics.micro_f1
+            )
+
+            if val_loss < best_val_loss:
+                patience = self.patience
+                best_val_loss = val_loss
+                logger.info("** Validation improved, evaluating test data **")
+                test_preds, segments, valid_len, test_loss = self.eval(self.test_dataloader)
+                self.segments_to_file(segments, os.path.join(self.output_path, "predictions.txt"))
+                test_metrics = compute_nested_metrics(segments, self.test_dataloader.dataset.transform.vocab.tags[1:])
+
+                epoch_summary_loss["test_loss"] = test_loss
+                epoch_summary_metrics["test_micro_f1"] = test_metrics.micro_f1
+                epoch_summary_metrics["test_precision"] = test_metrics.precision
+                epoch_summary_metrics["test_recall"] = test_metrics.recall
+
+                logger.info(
+                    f"Epoch %d | Timestep %d | Test Loss %f | F1 %f",
+                    epoch_index,
+                    self.current_timestep,
+                    test_loss,
+                    test_metrics.micro_f1
+                )
+
+                self.save()
+            else:
+                patience -= 1
+
+            # No improvements, terminating early
+            if patience == 0:
+                logger.info("Early termination triggered")
+                break
+
+            self.summary_writer.add_scalars("Loss", epoch_summary_loss, global_step=self.current_timestep)
+            self.summary_writer.add_scalars("Metrics", epoch_summary_metrics, global_step=self.current_timestep)
+
+    def tag(self, dataloader, is_train=True):
+        """
+        Given a dataloader containing segments, predict the tags
+        :param dataloader: torch.utils.data.DataLoader
+        :param is_train: boolean - True for training model, False for evaluation
+        :return: Iterator
+                    subwords (B x T x NUM_LABELS)- torch.Tensor - BERT subword ID
+                    gold_tags (B x T x NUM_LABELS) - torch.Tensor - ground truth tags IDs
+                    tokens - List[arabiner.data.dataset.Token] - list of tokens
+                    valid_len (B x 1) - int - valiud length of each sequence
+                    logits (B x T x NUM_LABELS) - logits for each token and each tag
+        """
+        for subwords, gold_tags, tokens, mask, valid_len in dataloader:
+            self.model.train(is_train)
+
+            if torch.cuda.is_available():
+                subwords = subwords.cuda()
+                gold_tags = gold_tags.cuda()
+
+            if is_train:
+                self.optimizer.zero_grad()
+                logits = self.model(subwords)
+            else:
+                with torch.no_grad():
+                    logits = self.model(subwords)
+
+            yield subwords, gold_tags, tokens, valid_len, logits
+
+    def eval(self, dataloader):
+        golds, preds, segments, valid_lens = list(), list(), list(), list()
+        num_labels = [len(v) for v in dataloader.dataset.vocab.tags[1:]]
+        loss = 0
+
+        for _, gold_tags, tokens, valid_len, logits in self.tag(
+            dataloader, is_train=False
+        ):
+            losses = [self.loss(logits[:, :, i, 0:l].view(-1, logits[:, :, i, 0:l].shape[-1]),
+                                torch.reshape(gold_tags[:, i, :], (-1,)).long())
+                      for i, l in enumerate(num_labels)]
+            loss += sum(losses)
+            preds += torch.argmax(logits, dim=3)
+            segments += tokens
+            valid_lens += list(valid_len)
+
+        loss /= len(dataloader)
+
+        # Update segments, attach predicted tags to each token
+        segments = self.to_segments(segments, preds, valid_lens, dataloader.dataset.vocab)
+
+        return preds, segments, valid_lens, loss
+
+    def infer(self, dataloader):
+        golds, preds, segments, valid_lens = list(), list(), list(), list()
+
+        for _, gold_tags, tokens, valid_len, logits in self.tag(
+            dataloader, is_train=False
+        ):
+            preds += torch.argmax(logits, dim=3)
+            segments += tokens
+            valid_lens += list(valid_len)
+
+        segments = self.to_segments(segments, preds, valid_lens, dataloader.dataset.vocab)
+        return segments
+
+    def to_segments(self, segments, preds, valid_lens, vocab):
+        if vocab is None:
+            vocab = self.vocab
+
+        tagged_segments = list()
+        tokens_stoi = vocab.tokens.get_stoi()
+        unk_id = tokens_stoi["UNK"]
+
+        for segment, pred, valid_len in zip(segments, preds, valid_lens):
+            # First, the token at 0th index [CLS] and token at nth index [SEP]
+            # Combine the tokens with their corresponding predictions
+            segment_pred = zip(segment[1:valid_len-1], pred[1:valid_len-1])
+
+            # Ignore the sub-tokens/subwords, which are identified with text being UNK
+            segment_pred = list(filter(lambda t: tokens_stoi[t[0].text] != unk_id, segment_pred))
+
+            # Attach the predicted tags to each token
+            list(map(lambda t: setattr(t[0], 'pred_tag', [{"tag": vocab.get_itos()[tag_id]}
+                                                     for tag_id, vocab in zip(t[1].int().tolist(), vocab.tags[1:])]), segment_pred))
+
+            # We are only interested in the tagged tokens, we do no longer need raw model predictions
+            tagged_segment = [t for t, _ in segment_pred]
+            tagged_segments.append(tagged_segment)
+
+        return tagged_segments