SinaTools 0.1.1__py2.py3-none-any.whl

nlptools/arabiner/trainers/BertTrainer.py

@@ -0,0 +1,163 @@
+import os
+import logging
+import torch
+import numpy as np
+from nlptools.arabiner.trainers import BaseTrainer
+from nlptools.arabiner.utils.metrics import compute_single_label_metrics
+    
+logger = logging.getLogger(__name__)
+
+
+class BertTrainer(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)
+        patience = self.patience
+
+        for epoch_index in range(self.max_epochs):
+            self.current_epoch = epoch_index
+            train_loss = 0
+
+            for batch_index, (_, gold_tags, _, _, logits) in enumerate(self.tag(
+                self.train_dataloader, is_train=True
+            ), 1):
+                self.current_timestep += 1
+                batch_loss = self.loss(logits.view(-1, logits.shape[-1]), gold_tags.view(-1))
+                batch_loss.backward()
+
+                # Avoid exploding gradient by doing gradient clipping
+                torch.nn.utils.clip_grad_norm_(self.model.parameters(), self.clip)
+
+                self.optimizer.step()
+                self.scheduler.step()
+                train_loss += batch_loss.item()
+
+                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.item()
+                    )
+
+            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_single_label_metrics(segments)
+
+            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_single_label_metrics(segments)
+
+                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 eval(self, dataloader):
+        golds, preds, segments, valid_lens = list(), list(), list(), list()
+        loss = 0
+
+        for _, gold_tags, tokens, valid_len, logits in self.tag(
+            dataloader, is_train=False
+        ):
+            loss += self.loss(logits.view(-1, logits.shape[-1]), gold_tags.view(-1))
+            preds += torch.argmax(logits, dim=2).detach().cpu().numpy().tolist()
+            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.item()
+
+    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=2).detach().cpu().numpy().tolist()
+            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()
+        tags_itos = vocab.tags[0].get_itos()
+        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": tags_itos[t[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

nlptools/arabiner/trainers/__init__.py

@@ -0,0 +1,3 @@
+from nlptools.arabiner.trainers.BaseTrainer import BaseTrainer
+from nlptools.arabiner.trainers.BertTrainer import BertTrainer
+from nlptools.arabiner.trainers.BertNestedTrainer import BertNestedTrainer

nlptools/arabiner/utils/data.py

@@ -0,0 +1,124 @@
+from torch.utils.data import DataLoader
+from torchtext.vocab import vocab
+from collections import Counter, namedtuple
+import logging
+import re
+import itertools
+from nlptools.arabiner.utils.helpers import load_object
+from nlptools.arabiner.data.datasets import Token
+from nlptools.morphology.tokenizers_words import simple_word_tokenize
+
+logger = logging.getLogger(__name__)
+
+
+def conll_to_segments(filename):
+    """
+    Convert CoNLL files to segments. This return list of segments and each segment is
+    a list of tuples (token, tag)
+    :param filename: Path
+    :return: list[[tuple]] - [[(token, tag), (token, tag), ...], [(token, tag), ...]]
+    """
+    segments, segment = list(), list()
+
+    with open(filename, "r") as fh:
+        for token in fh.read().splitlines():
+            if not token.strip():
+                segments.append(segment)
+                segment = list()
+            else:
+                parts = token.split()
+                token = Token(text=parts[0], gold_tag=parts[1:])
+                segment.append(token)
+
+        segments.append(segment)
+
+    return segments
+
+
+def parse_conll_files(data_paths):
+    """
+    Parse CoNLL formatted files and return list of segments for each file and index
+    the vocabs and tags across all data_paths
+    :param data_paths: tuple(Path) - tuple of filenames
+    :return: tuple( [[(token, tag), ...], [(token, tag), ...]], -> segments for data_paths[i]
+                    [[(token, tag), ...], [(token, tag), ...]], -> segments for data_paths[i+1],
+                    ...
+                  )
+             List of segments for each dataset and each segment has list of (tokens, tags)
+    """
+    vocabs = namedtuple("Vocab", ["tags", "tokens"])
+    datasets, tags, tokens = list(), list(), list()
+
+    for data_path in data_paths:
+        dataset = conll_to_segments(data_path)
+        datasets.append(dataset)
+        tokens += [token.text for segment in dataset for token in segment]
+        tags += [token.gold_tag for segment in dataset for token in segment]
+
+    # Flatten list of tags
+    tags = list(itertools.chain(*tags))
+
+    # Generate vocabs for tags and tokens
+    tag_vocabs = tag_vocab_by_type(tags)
+    tag_vocabs.insert(0, vocab(Counter(tags)))
+    vocabs = vocabs(tokens=vocab(Counter(tokens), specials=["UNK"]), tags=tag_vocabs)
+    return tuple(datasets), vocabs
+
+
+def tag_vocab_by_type(tags):
+    vocabs = list()
+    c = Counter(tags)
+    tag_names = c.keys()
+    tag_types = sorted(list(set([tag.split("-", 1)[1] for tag in tag_names if "-" in tag])))
+
+    for tag_type in tag_types:
+        r = re.compile(".*-" + tag_type)
+        t = list(filter(r.match, tags)) + ["O"]
+        vocabs.append(vocab(Counter(t), specials=["<pad>"]))
+
+    return vocabs
+
+
+def text2segments(text):
+    """
+    Convert text to a datasets and index the tokens
+    """
+    #dataset = [[Token(text=token, gold_tag=["O"]) for token in text.split()]]
+    list_of_tokens = simple_word_tokenize(text)
+    dataset = [[Token(text=token, gold_tag=["O"]) for token in list_of_tokens]]
+    tokens = [token.text for segment in dataset for token in segment]
+
+    # Generate vocabs for the tokens
+    segment_vocab = vocab(Counter(tokens), specials=["UNK"])
+    return dataset, segment_vocab
+
+
+def get_dataloaders(
+    datasets, vocab, data_config, batch_size=32, num_workers=0, shuffle=(True, False, False)
+):
+    """
+    From the datasets generate the dataloaders
+    :param datasets: list - list of the datasets, list of list of segments and tokens
+    :param batch_size: int
+    :param num_workers: int
+    :param shuffle: boolean - to shuffle the data or not
+    :return: List[torch.utils.data.DataLoader]
+    """
+    dataloaders = list()
+
+    for i, examples in enumerate(datasets):
+        data_config["kwargs"].update({"examples": examples, "vocab": vocab})
+        dataset = load_object("nlptools."+data_config["fn"], data_config["kwargs"])
+
+        dataloader = DataLoader(
+            dataset=dataset,
+            shuffle=shuffle[i],
+            batch_size=batch_size,
+            num_workers=num_workers,
+            collate_fn=dataset.collate_fn,
+        )
+
+        logger.info("%s batches found", len(dataloader))
+        dataloaders.append(dataloader)
+
+    return dataloaders

nlptools/arabiner/utils/helpers.py

@@ -0,0 +1,151 @@
+import os
+import sys
+import logging
+import importlib
+import shutil
+import torch
+import pickle
+import json
+import random
+import numpy as np
+from argparse import Namespace
+
+
+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):
+#    """
+#    Load objects dynamically given the object name and its arguments
+#    :param name: str - object name, class name or function name
+#    :param kwargs: dict - keyword arguments
+#    :return: object
+#    """
+#    object_module, object_name = name.rsplit(".", 1)
+#    object_module = importlib.import_module(object_module)
+#    fn = getattr(object_module, object_name)(**kwargs)
+#    return fn
+
+def load_object(name, kwargs):
+    #print("Iterate: ", name, kwargs)
+    try:
+        #print("name: ", name)
+        object_module, object_name = name.rsplit(".", 1)
+        #print("object_module: ", object_module)
+        #print("object_name: ", object_name)
+        #if object_module != "nlptools.arabiner.nn":
+        object_module = importlib.import_module(object_module)
+        #print("object_module 2: ", object_module)
+        obj = getattr(object_module, object_name)
+        #print("obj: ", obj)
+        if callable(obj):  # Check if the object is callable (class or function)
+            fn = obj(**kwargs)
+            return fn
+        else:
+            raise TypeError(f"{name} is not a callable object.")
+    except (ImportError, ModuleNotFoundError) as e:
+        # Handle import errors
+        print(f"Error importing module: {e}")
+    except AttributeError as e:
+        # Handle attribute errors (e.g., object not found in module)
+        print(f"Attribute error: {e}")
+    except Exception as e:
+        # Handle other exceptions
+        print(f"An error occurred: {e}")
+    
+    #print("Loaded object:", name)
+    return None  # Return None in case of any error
+
+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 load_checkpoint(path):
+    """
+    Load model given the model path
+    :param model_path: str - path to model
+    :return: tagger - arabiner.trainers.BaseTrainer - the tagger model
+             vocab - torchtext.vocab.Vocab - indexed tags
+             train_config - argparse.Namespace - training configurations
+    """
+    _path = os.path.join(path, "tag_vocab.pkl")
+    #print('2',_path)
+    with open(_path, "rb") as fh:
+        tag_vocab = pickle.load(fh)
+
+    # Load train configurations from checkpoint
+    train_config = Namespace()
+    args_path = os.path.join(path, "args.json")
+    #print('3', args_path)
+    with open(args_path, "r") as fh:
+        train_config.__dict__ = json.load(fh)
+
+    # Initialize the loss function, not used for inference, but evaluation
+    loss = load_object(train_config.loss["fn"], train_config.loss["kwargs"])
+    #print('4')
+    # Load BERT tagger
+    model = load_object("nlptools."+train_config.network_config["fn"], train_config.network_config["kwargs"])
+    model = torch.nn.DataParallel(model)
+    #print('5')
+    if torch.cuda.is_available():
+        model = model.cuda()
+
+    # Update arguments for the tagger
+    # Attach the model, loss (used for evaluations cases)
+    train_config.trainer_config["kwargs"]["model"] = model
+    train_config.trainer_config["kwargs"]["loss"] = loss
+    #print('6') 
+    tagger = load_object("nlptools."+train_config.trainer_config["fn"], train_config.trainer_config["kwargs"])
+    tagger.load(os.path.join(path, "checkpoints"))
+    #print('7') 
+    return tagger, tag_vocab, train_config
+
+
+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

nlptools/arabiner/utils/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)