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

sinatools/ner/data/datasets.py

@@ -0,0 +1,146 @@
+import logging
+import torch
+from torch.utils.data import Dataset
+from torch.nn.utils.rnn import pad_sequence
+from sinatools.ner.data.transforms import (
+    BertSeqTransform,
+    NestedTagsTransform
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Token:
+    def __init__(self, text=None, pred_tag=None, gold_tag=None):
+        """
+        Token object to hold token attributes
+        :param text: str
+        :param pred_tag: str
+        :param gold_tag: str
+        """
+        self.text = text
+        self.gold_tag = gold_tag
+        self.pred_tag = pred_tag
+        self.subwords = None
+
+    @property
+    def subwords(self):
+        return self._subwords
+
+    @subwords.setter
+    def subwords(self, value):
+        self._subwords = value
+
+    def __str__(self):
+        """
+        Token text representation
+        :return: str
+        """
+        gold_tags = "|".join(self.gold_tag)
+        pred_tags = "|".join([pred_tag["tag"] for pred_tag in self.pred_tag])
+
+        if self.gold_tag:
+            r = f"{self.text}\t{gold_tags}\t{pred_tags}"
+        else:
+            r = f"{self.text}\t{pred_tags}"
+
+        return r
+
+
+class DefaultDataset(Dataset):
+    def __init__(
+        self,
+        examples=None,
+        vocab=None,
+        bert_model="aubmindlab/bert-base-arabertv2",
+        max_seq_len=512,
+    ):
+        """
+        The dataset that used to transform the segments into training data
+        :param examples: list[[tuple]] - [[(token, tag), (token, tag), ...], [(token, tag), ...]]
+                         You can get generate examples from -- arabiner.data.dataset.parse_conll_files
+        :param vocab: vocab object containing indexed tags and tokens
+        :param bert_model: str - BERT model
+        :param: int - maximum sequence length
+        """
+        self.transform = BertSeqTransform(bert_model, vocab, max_seq_len=max_seq_len)
+        self.examples = examples
+        self.vocab = vocab
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        subwords, tags, tokens, valid_len = self.transform(self.examples[item])
+        return subwords, tags, tokens, valid_len
+
+    def collate_fn(self, batch):
+        """
+        Collate function that is called when the batch is called by the trainer
+        :param batch: Dataloader batch
+        :return: Same output as the __getitem__ function
+        """
+        subwords, tags, tokens, valid_len = zip(*batch)
+
+        # Pad sequences in this batch
+        # subwords and tokens are padded with zeros
+        # tags are padding with the index of the O tag
+        subwords = pad_sequence(subwords, batch_first=True, padding_value=0)
+        tags = pad_sequence(
+            tags, batch_first=True, padding_value=self.vocab.tags[0].get_stoi()["O"]
+        )
+        return subwords, tags, tokens, valid_len
+
+
+class NestedTagsDataset(Dataset):
+    def __init__(
+        self,
+        examples=None,
+        vocab=None,
+        bert_model="aubmindlab/bert-base-arabertv2",
+        max_seq_len=512,
+    ):
+        """
+        The dataset that used to transform the segments into training data
+        :param examples: list[[tuple]] - [[(token, tag), (token, tag), ...], [(token, tag), ...]]
+                         You can get generate examples from -- arabiner.data.dataset.parse_conll_files
+        :param vocab: vocab object containing indexed tags and tokens
+        :param bert_model: str - BERT model
+        :param: int - maximum sequence length
+        """
+        self.transform = NestedTagsTransform(
+            bert_model, vocab, max_seq_len=max_seq_len
+        )
+        self.examples = examples
+        self.vocab = vocab
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        subwords, tags, tokens, masks, valid_len = self.transform(self.examples[item])
+        return subwords, tags, tokens, masks, valid_len
+
+    def collate_fn(self, batch):
+        """
+        Collate function that is called when the batch is called by the trainer
+        :param batch: Dataloader batch
+        :return: Same output as the __getitem__ function
+        """
+        subwords, tags, tokens, masks, valid_len = zip(*batch)
+
+        # Pad sequences in this batch
+        # subwords and tokens are padded with zeros
+        # tags are padding with the index of the O tag
+        subwords = pad_sequence(subwords, batch_first=True, padding_value=0)
+
+        masks = [torch.nn.ConstantPad1d((0, subwords.shape[-1] - tag.shape[-1]), 0)(mask)
+                 for tag, mask in zip(tags, masks)]
+        masks = torch.cat(masks)
+
+        # Pad the tags, do the padding for each tag type
+        tags = [torch.nn.ConstantPad1d((0, subwords.shape[-1] - tag.shape[-1]), vocab.get_stoi()["<pad>"])(tag)
+                for tag, vocab in zip(tags, self.vocab.tags[1:])]
+        tags = torch.cat(tags)
+
+        return subwords, tags, tokens, masks, valid_len

sinatools/ner/data/transforms.py

@@ -0,0 +1,118 @@
+import torch
+from transformers import BertTokenizer
+from functools import partial
+import re
+import itertools
+from sinatools.ner.data import datasets
+class BertSeqTransform:
+    def __init__(self, bert_model, vocab, max_seq_len=512):
+        self.tokenizer = BertTokenizer.from_pretrained(bert_model)
+        self.encoder = partial(
+            self.tokenizer.encode,
+            max_length=max_seq_len,
+            truncation=True,
+        )
+        self.max_seq_len = max_seq_len
+        self.vocab = vocab
+
+    def __call__(self, segment):
+        subwords, tags, tokens = list(), list(), list()
+        unk_token = datasets.Token(text="UNK")
+
+        for token in segment:
+            token_subwords = self.encoder(token.text)[1:-1]
+            subwords += token_subwords
+            tags += [self.vocab.tags[0].get_stoi()[token.gold_tag[0]]] + [self.vocab.tags[0].get_stoi()["O"]] * (len(token_subwords) - 1)
+            tokens += [token] + [unk_token] * (len(token_subwords) - 1)
+
+        # Truncate to max_seq_len
+        if len(subwords) > self.max_seq_len - 2:
+            text = " ".join([t.text for t in tokens if t.text != "UNK"])
+          
+            subwords = subwords[:self.max_seq_len - 2]
+            tags = tags[:self.max_seq_len - 2]
+            tokens = tokens[:self.max_seq_len - 2]
+
+        subwords.insert(0, self.tokenizer.cls_token_id)
+        subwords.append(self.tokenizer.sep_token_id)
+
+        tags.insert(0, self.vocab.tags[0].get_stoi()["O"])
+        tags.append(self.vocab.tags[0].get_stoi()["O"])
+
+        tokens.insert(0, unk_token)
+        tokens.append(unk_token)
+
+        return torch.LongTensor(subwords), torch.LongTensor(tags), tokens, len(tokens)
+
+
+class NestedTagsTransform:
+    def __init__(self, bert_model, vocab, max_seq_len=512):
+        self.tokenizer = BertTokenizer.from_pretrained(bert_model)
+        self.encoder = partial(
+            self.tokenizer.encode,
+            max_length=max_seq_len,
+            truncation=True,
+        )
+        self.max_seq_len = max_seq_len
+        self.vocab = vocab
+
+    def __call__(self, segment):
+        tags, tokens, subwords = list(), list(), list()
+        unk_token = datasets.Token(text="UNK")
+
+        # Encode each token and get its subwords and IDs
+        for token in segment:
+            token.subwords = self.encoder(token.text)[1:-1]
+            subwords += token.subwords
+            tokens += [token] + [unk_token] * (len(token.subwords ) - 1)
+
+        # Construct the labels for each tag type
+        # The sequence will have a list of tags for each type
+        # The final tags for a sequence is a matrix NUM_TAG_TYPES x SEQ_LEN
+        # Example:
+        #   [
+        #       [O,     O,     B-PERS, I-PERS, O, O, O]
+        #       [B-ORG, I-ORG, O,      O,      O, O, O]
+        #       [O,     O,     O,      O,      O, O, B-GPE]
+        #   ]
+        for vocab in self.vocab.tags[1:]:
+            vocab_tags = "|".join([t for t in vocab.get_itos() if "-" in t])
+            r = re.compile(vocab_tags)
+
+            # This is really messy
+            # For a given token we find a matching tag_name, BUT we might find
+            # multiple matches (i.e. a token can be labeled B-ORG and I-ORG) in this
+            # case we get only the first tag as we do not have overlapping of same type
+            single_type_tags = [[(list(filter(r.match, token.gold_tag))
+                                or ["O"])[0]] + ["O"] * (len(token.subwords) - 1)
+                                for token in segment]
+            single_type_tags = list(itertools.chain(*single_type_tags))
+            tags.append([vocab.get_stoi()[tag] for tag in single_type_tags])
+
+        # Truncate to max_seq_len
+        if len(subwords) > self.max_seq_len - 2:
+            text = " ".join([t.text for t in tokens if t.text != "UNK"])
+          
+            subwords = subwords[:self.max_seq_len - 2]
+            tags = [t[:self.max_seq_len - 2] for t in tags]
+            tokens = tokens[:self.max_seq_len - 2]
+
+        # Add dummy token at the start end of sequence
+        tokens.insert(0, unk_token)
+        tokens.append(unk_token)
+
+        # Add CLS and SEP at start end of subwords
+        subwords.insert(0, self.tokenizer.cls_token_id)
+        subwords.append(self.tokenizer.sep_token_id)
+        subwords = torch.LongTensor(subwords)
+
+        # Add "O" tags for the first and last subwords
+        tags = torch.Tensor(tags)
+        tags = torch.column_stack((
+            torch.Tensor([vocab.get_stoi()["O"] for vocab in self.vocab.tags[1:]]),
+            tags,
+            torch.Tensor([vocab.get_stoi()["O"] for vocab in self.vocab.tags[1:]]),
+        )).unsqueeze(0)
+
+        mask = torch.ones_like(tags)
+        return subwords, tags, tokens, mask, len(tokens)

sinatools/ner/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 sinatools.ner.helpers import load_object
+from sinatools.ner.datasets import Token
+from sinatools.utils.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("sinatools."+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

sinatools/ner/data_format.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 sinatools.ner.helpers import load_object
+from sinatools.ner.datasets import Token
+from sinatools.utils.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(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

sinatools/ner/datasets.py

@@ -0,0 +1,146 @@
+import logging
+import torch
+from torch.utils.data import Dataset
+from torch.nn.utils.rnn import pad_sequence
+from sinatools.ner.transforms import (
+    BertSeqTransform,
+    NestedTagsTransform
+)
+
+logger = logging.getLogger(__name__)
+
+
+class Token:
+    def __init__(self, text=None, pred_tag=None, gold_tag=None):
+        """
+        Token object to hold token attributes
+        :param text: str
+        :param pred_tag: str
+        :param gold_tag: str
+        """
+        self.text = text
+        self.gold_tag = gold_tag
+        self.pred_tag = pred_tag
+        self.subwords = None
+
+    @property
+    def subwords(self):
+        return self._subwords
+
+    @subwords.setter
+    def subwords(self, value):
+        self._subwords = value
+
+    def __str__(self):
+        """
+        Token text representation
+        :return: str
+        """
+        gold_tags = "|".join(self.gold_tag)
+        pred_tags = "|".join([pred_tag["tag"] for pred_tag in self.pred_tag])
+
+        if self.gold_tag:
+            r = f"{self.text}\t{gold_tags}\t{pred_tags}"
+        else:
+            r = f"{self.text}\t{pred_tags}"
+
+        return r
+
+
+class DefaultDataset(Dataset):
+    def __init__(
+        self,
+        examples=None,
+        vocab=None,
+        bert_model="aubmindlab/bert-base-arabertv2",
+        max_seq_len=512,
+    ):
+        """
+        The dataset that used to transform the segments into training data
+        :param examples: list[[tuple]] - [[(token, tag), (token, tag), ...], [(token, tag), ...]]
+                         You can get generate examples from -- arabiner.data.dataset.parse_conll_files
+        :param vocab: vocab object containing indexed tags and tokens
+        :param bert_model: str - BERT model
+        :param: int - maximum sequence length
+        """
+        self.transform = BertSeqTransform(bert_model, vocab, max_seq_len=max_seq_len)
+        self.examples = examples
+        self.vocab = vocab
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        subwords, tags, tokens, valid_len = self.transform(self.examples[item])
+        return subwords, tags, tokens, valid_len
+
+    def collate_fn(self, batch):
+        """
+        Collate function that is called when the batch is called by the trainer
+        :param batch: Dataloader batch
+        :return: Same output as the __getitem__ function
+        """
+        subwords, tags, tokens, valid_len = zip(*batch)
+
+        # Pad sequences in this batch
+        # subwords and tokens are padded with zeros
+        # tags are padding with the index of the O tag
+        subwords = pad_sequence(subwords, batch_first=True, padding_value=0)
+        tags = pad_sequence(
+            tags, batch_first=True, padding_value=self.vocab.tags[0].get_stoi()["O"]
+        )
+        return subwords, tags, tokens, valid_len
+
+
+class NestedTagsDataset(Dataset):
+    def __init__(
+        self,
+        examples=None,
+        vocab=None,
+        bert_model="aubmindlab/bert-base-arabertv2",
+        max_seq_len=512,
+    ):
+        """
+        The dataset that used to transform the segments into training data
+        :param examples: list[[tuple]] - [[(token, tag), (token, tag), ...], [(token, tag), ...]]
+                         You can get generate examples from -- arabiner.data.dataset.parse_conll_files
+        :param vocab: vocab object containing indexed tags and tokens
+        :param bert_model: str - BERT model
+        :param: int - maximum sequence length
+        """
+        self.transform = NestedTagsTransform(
+            bert_model, vocab, max_seq_len=max_seq_len
+        )
+        self.examples = examples
+        self.vocab = vocab
+
+    def __len__(self):
+        return len(self.examples)
+
+    def __getitem__(self, item):
+        subwords, tags, tokens, masks, valid_len = self.transform(self.examples[item])
+        return subwords, tags, tokens, masks, valid_len
+
+    def collate_fn(self, batch):
+        """
+        Collate function that is called when the batch is called by the trainer
+        :param batch: Dataloader batch
+        :return: Same output as the __getitem__ function
+        """
+        subwords, tags, tokens, masks, valid_len = zip(*batch)
+
+        # Pad sequences in this batch
+        # subwords and tokens are padded with zeros
+        # tags are padding with the index of the O tag
+        subwords = pad_sequence(subwords, batch_first=True, padding_value=0)
+
+        masks = [torch.nn.ConstantPad1d((0, subwords.shape[-1] - tag.shape[-1]), 0)(mask)
+                 for tag, mask in zip(tags, masks)]
+        masks = torch.cat(masks)
+
+        # Pad the tags, do the padding for each tag type
+        tags = [torch.nn.ConstantPad1d((0, subwords.shape[-1] - tag.shape[-1]), vocab.get_stoi()["<pad>"])(tag)
+                for tag, vocab in zip(tags, self.vocab.tags[1:])]
+        tags = torch.cat(tags)
+
+        return subwords, tags, tokens, masks, valid_len