SinaTools 0.1.11__py2.py3-none-any.whl → 0.1.13__py2.py3-none-any.whl

sinatools/ner/entity_extractor.py

@@ -1,51 +1,51 @@
 import os
 from collections import namedtuple
-from sinatools.ner.utils.helpers import load_checkpoint
-from sinatools.ner.utils.data import get_dataloaders, text2segments
-from sinatools.DataDownload import downloader
-from . import tag_vocab, train_config, tagger
+from sinatools.ner.data_format import get_dataloaders, text2segments
+from . import tagger, tag_vocab, train_config
 
-def ner(text, batch_size=32):
+def extract(text, batch_size=32):
     """
-    This method takes a text as input, and a batch size, then performs named entity recognition (NER) on the input text and returns a list of tagged mentions.
+    This method processes an input text and returns named entites for each token within the text, based on the specified batch size. As follows:
 
     Args:
-        text (str): The input text to perform NER on.
+        text (:obj:`str`): The Arabic text to be tagged.
         batch_size (int, optional): Batch size for inference. Default is 32.
 
     Returns:
-        list: A list of lists containing token and label pairs for each segment.
-              Each inner list has the format ['token', 'label1 label2 ...'].
+        list (:obj:`list`): A list of JSON objects, where each JSON could be contains:
+        token: The token from the original text.
+        NER tag: The label pairs for each segment.
+
     **Example:**
 
      .. highlight:: python
      .. code-block:: python
 
-            from sinatools.arabiner.bin import infer
-            infer.ner('ذهب محمد الى جامعة بيرزيت')
-
-            #the output   
-            [['ذهب', 'O'],
-            ['محمد', 'B-PERS'],
-            ['الى', 'O'],
-            ['جامعة', 'B-ORG'],
-            ['بيرزيت', 'B-GPE I-ORG']]
+        from sinatools.ner.entity_extractor import extract
+        extract('ذهب محمد إلى جامعة بيرزيت')
+        [{
+            "word":"ذهب",
+            "tags":"O"
+          },{
+            "word":"محمد",
+            "tags":"B-PERS"
+          },{
+            "word":"إلى",
+            "tags":"O"
+          },{
+            "word":"جامعة",
+            "tags":"B-ORG"
+          },{
+            "word":"بيرزيت",
+            "tags":"B-GPE I-ORG"
+        }]
     """    
-    # Load tagger
-    # filename = 'Wj27012000.tar'
-    # path =downloader.get_appdatadir()
-    # model_path = os.path.join(path, filename)
-    # print('1',model_path)
-    # tagger, tag_vocab, train_config = load_checkpoint(model_path)
     
-
-    # Convert text to a tagger dataset and index the tokens in args.text
     dataset, token_vocab = text2segments(text)
 
     vocabs = namedtuple("Vocab", ["tags", "tokens"])
     vocab = vocabs(tokens=token_vocab, tags=tag_vocab)
 
-    # From the datasets generate the dataloaders
     dataloader = get_dataloaders(
         (dataset,),
         vocab,
@@ -54,39 +54,19 @@ def ner(text, batch_size=32):
         shuffle=(False,),
     )[0]
 
-    # Perform inference on the text and get back the tagged segments
+
     segments = tagger.infer(dataloader)
     segments_lists = []
-    # Print results
+    
     for segment in segments:
         for token in segment:
-            segments_list = []
-            segments_list.append(token.text)
-            #print([t['tag'] for t in token.pred_tag])
+            segments_list = {}
+            segments_list["word"] = token.text
             list_of_tags = [t['tag'] for t in token.pred_tag]
             list_of_tags = [i for i in list_of_tags if i not in('O',' ','')]
-            #print(list_of_tags)
             if list_of_tags == []:
-               segments_list.append(' '.join(['O']))
+               segments_list["tag"] = ' '.join(['O'])
             else:
-               segments_list.append(' '.join(list_of_tags))
-            segments_lists.append(segments_list)         
+               segments_list["tag"] = ' '.join(list_of_tags)
+            segments_lists.append(segments_list)  
     return segments_lists
-
-#Print results
-    # for segment in segments:
-    #     s = [
-    #         (token.text, token.pred_tag[0]['tag'])
-    #         for token in segment
-    #         if token.pred_tag[0]['tag'] != 'O'
-    #     ]
-    #     print(", ".join([f"({token}, {tag})" for token, tag in s]))
-
-def extract_tags(text):
-    tags = []
-    tokens = text.split()
-    for token in tokens:
-        tag = token.split("(")[-1].split(")")[0]
-        if tag != "O":
-            tags.append(tag)
-    return " ".join(tags)

sinatools/ner/helpers.py

@@ -0,0 +1,86 @@
+import os
+import sys
+import logging
+import importlib
+import shutil
+import torch
+import random
+import numpy as np
+
+
+def logging_config(log_file=None):
+    """
+    Initialize custom logger
+    :param log_file: str - path to log file, full path
+    :return: None
+    """
+    handlers = [logging.StreamHandler(sys.stdout)]
+
+    if log_file:
+        handlers.append(logging.FileHandler(log_file, "w", "utf-8"))
+        print("Logging to {}".format(log_file))
+
+    logging.basicConfig(
+        level=logging.INFO,
+        handlers=handlers,
+        format="%(levelname)s\t%(name)s\t%(asctime)s\t%(message)s",
+        datefmt="%a, %d %b %Y %H:%M:%S",
+        force=True
+    )
+
+
+def load_object(name, kwargs):
+
+    try:
+        object_module, object_name = name.rsplit(".", 1)
+        object_module = importlib.import_module(object_module)
+        obj = getattr(object_module, object_name)
+        if callable(obj): 
+            fn = obj(**kwargs)
+            return fn
+        else:
+            raise TypeError(f"{name} is not a callable object.")
+    except (ImportError, ModuleNotFoundError) as e:
+        print(f"Error importing module: {e}")
+    except AttributeError as e:
+        print(f"Attribute error: {e}")
+    except Exception as e:
+        print(f"An error occurred: {e}")
+    
+    return None 
+
+def make_output_dirs(path, subdirs=[], overwrite=True):
+    """
+    Create root directory and any other sub-directories
+    :param path: str - root directory
+    :param subdirs: List[str] - list of sub-directories
+    :param overwrite: boolean - to overwrite the directory or not
+    :return: None
+    """
+    if overwrite:
+        shutil.rmtree(path, ignore_errors=True)
+
+    os.makedirs(path)
+
+    for subdir in subdirs:
+        os.makedirs(os.path.join(path, subdir))
+
+
+
+def set_seed(seed):
+    """
+    Set the seed for random intialization and set
+    CUDANN parameters to ensure determmihstic results across
+    multiple runs with the same seed
+
+    :param seed: int
+    """
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.enabled = False

sinatools/ner/metrics.py

@@ -0,0 +1,69 @@
+from seqeval.metrics import (
+    classification_report,
+    precision_score,
+    recall_score,
+    f1_score,
+    accuracy_score,
+)
+from seqeval.scheme import IOB2
+from types import SimpleNamespace
+import logging
+import re
+
+logger = logging.getLogger(__name__)
+
+
+def compute_nested_metrics(segments, vocabs):
+    """
+    Compute metrics for nested NER
+    :param segments: List[List[arabiner.data.dataset.Token]] - list of segments
+    :return: metrics - SimpleNamespace - F1/micro/macro/weights, recall, precision, accuracy
+    """
+    y, y_hat = list(), list()
+
+    # We duplicate the dataset N times, where N is the number of entity types
+    # For each copy, we create y and y_hat
+    # Example: first copy, will create pairs of ground truth and predicted labels for entity type GPE
+    #          another copy will create pairs for LOC, etc.
+    for i, vocab in enumerate(vocabs):
+        vocab_tags = [tag for tag in vocab.get_itos() if "-" in tag]
+        r = re.compile("|".join(vocab_tags))
+
+        y += [[(list(filter(r.match, token.gold_tag)) or ["O"])[0] for token in segment] for segment in segments]
+        y_hat += [[token.pred_tag[i]["tag"] for token in segment] for segment in segments]
+
+    logging.info("\n" + classification_report(y, y_hat, scheme=IOB2, digits=4))
+
+    metrics = {
+        "micro_f1": f1_score(y, y_hat, average="micro", scheme=IOB2),
+        "macro_f1": f1_score(y, y_hat, average="macro", scheme=IOB2),
+        "weights_f1": f1_score(y, y_hat, average="weighted", scheme=IOB2),
+        "precision": precision_score(y, y_hat, scheme=IOB2),
+        "recall": recall_score(y, y_hat, scheme=IOB2),
+        "accuracy": accuracy_score(y, y_hat),
+    }
+
+    return SimpleNamespace(**metrics)
+
+
+def compute_single_label_metrics(segments):
+    """
+    Compute metrics for flat NER
+    :param segments: List[List[arabiner.data.dataset.Token]] - list of segments
+    :return: metrics - SimpleNamespace - F1/micro/macro/weights, recall, precision, accuracy
+    """
+    y = [[token.gold_tag[0] for token in segment] for segment in segments]
+    y_hat = [[token.pred_tag[0]["tag"] for token in segment] for segment in segments]
+
+    logging.info("\n" + classification_report(y, y_hat, scheme=IOB2))
+
+    metrics = {
+        "micro_f1": f1_score(y, y_hat, average="micro", scheme=IOB2),
+        "macro_f1": f1_score(y, y_hat, average="macro", scheme=IOB2),
+        "weights_f1": f1_score(y, y_hat, average="weighted", scheme=IOB2),
+        "precision": precision_score(y, y_hat, scheme=IOB2),
+        "recall": recall_score(y, y_hat, scheme=IOB2),
+        "accuracy": accuracy_score(y, y_hat),
+    }
+
+    return SimpleNamespace(**metrics)

sinatools/ner/nn/BaseModel.py

@@ -0,0 +1,22 @@
+from torch import nn
+from transformers import BertModel
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BaseModel(nn.Module):
+    def __init__(self,
+                 bert_model="aubmindlab/bert-base-arabertv2",
+                 num_labels=2,
+                 dropout=0.1,
+                 num_types=0):
+        super().__init__()
+
+        self.bert_model = bert_model
+        self.num_labels = num_labels
+        self.num_types = num_types
+        self.dropout = dropout
+
+        self.bert = BertModel.from_pretrained(bert_model)
+        self.dropout = nn.Dropout(dropout)

sinatools/ner/nn/BertNestedTagger.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+from sinatools.ner.nn import BaseModel
+
+
+class BertNestedTagger(BaseModel):
+    def __init__(self, **kwargs):
+        super(BertNestedTagger, self).__init__(**kwargs)
+
+        self.max_num_labels = max(self.num_labels)
+        classifiers = [nn.Linear(768, num_labels) for num_labels in self.num_labels]
+        self.classifiers = torch.nn.Sequential(*classifiers)
+
+    def forward(self, x):
+        y = self.bert(x)
+        y = self.dropout(y["last_hidden_state"])
+        output = list()
+
+        for i, classifier in enumerate(self.classifiers):
+            logits = classifier(y)
+
+            # Pad logits to allow Multi-GPU/DataParallel training to work
+            # We will truncate the padded dimensions when we compute the loss in the trainer
+            logits = torch.nn.ConstantPad1d((0, self.max_num_labels - logits.shape[-1]), 0)(logits)
+            output.append(logits)
+
+        # Return tensor of the shape B x T x L x C
+        # B: batch size
+        # T: sequence length
+        # L: number of tag types
+        # C: number of classes per tag type
+        output = torch.stack(output).permute((1, 2, 0, 3))
+        return output
+

sinatools/ner/nn/BertSeqTagger.py

@@ -0,0 +1,17 @@
+import torch.nn as nn
+from transformers import BertModel
+
+
+class BertSeqTagger(nn.Module):
+    def __init__(self, bert_model, num_labels=2, dropout=0.1):
+        super().__init__()
+
+        self.bert = BertModel.from_pretrained(bert_model)
+        self.dropout = nn.Dropout(dropout)
+        self.linear = nn.Linear(768, num_labels)
+
+    def forward(self, x):
+        y = self.bert(x)
+        y = self.dropout(y["last_hidden_state"])
+        logits = self.linear(y)
+        return logits

sinatools/ner/nn/__init__.py

@@ -0,0 +1,3 @@
+from sinatools.ner.nn.BaseModel import BaseModel
+from sinatools.ner.nn.BertSeqTagger import BertSeqTagger
+from sinatools.ner.nn.BertNestedTagger import BertNestedTagger

sinatools/ner/trainers/BaseTrainer.py

@@ -0,0 +1,117 @@
+import os
+import torch
+import logging
+import natsort
+import glob
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTrainer:
+    def __init__(
+        self,
+        model=None,
+        max_epochs=50,
+        optimizer=None,
+        scheduler=None,
+        loss=None,
+        train_dataloader=None,
+        val_dataloader=None,
+        test_dataloader=None,
+        log_interval=10,
+        summary_writer=None,
+        output_path=None,
+        clip=5,
+        patience=5
+    ):
+        self.model = model
+        self.max_epochs = max_epochs
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
+        self.test_dataloader = test_dataloader
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.loss = loss
+        self.log_interval = log_interval
+        self.summary_writer = summary_writer
+        self.output_path = output_path
+        self.current_timestep = 0
+        self.current_epoch = 0
+        self.clip = clip
+        self.patience = patience
+
+    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, 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 segments_to_file(self, segments, filename):
+        """
+        Write segments to file
+        :param segments: [List[arabiner.data.dataset.Token]] - list of list of tokens
+        :param filename: str - output filename
+        :return: None
+        """
+        with open(filename, "w") as fh:
+            results = "\n\n".join(["\n".join([t.__str__() for t in segment]) for segment in segments])
+            fh.write("Token\tGold Tag\tPredicted Tag\n")
+            fh.write(results)
+            logging.info("Predictions written to %s", filename)
+
+    def save(self):
+        """
+        Save model checkpoint
+        :return:
+        """
+        filename = os.path.join(
+            self.output_path,
+            "checkpoints",
+            "checkpoint_{}.pt".format(self.current_epoch),
+        )
+
+        checkpoint = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "epoch": self.current_epoch
+        }
+
+        logger.info("Saving checkpoint to %s", filename)
+        torch.save(checkpoint, filename)
+
+    def load(self, checkpoint_path):
+        """
+        Load model checkpoint
+        :param checkpoint_path: str - path/to/checkpoints
+        :return: None
+        """
+        checkpoint_path = natsort.natsorted(glob.glob(f"{checkpoint_path}/checkpoint_*.pt"))
+        checkpoint_path = checkpoint_path[-1]
+
+        logger.info("Loading checkpoint %s", checkpoint_path)
+
+        device = None if torch.cuda.is_available() else torch.device('cpu')
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        self.model.load_state_dict(checkpoint["model"])

sinatools/ner/trainers/BertNestedTrainer.py

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