BERT MLM Support

syfertext/data/iterators/bert_loader.py

@@ -0,0 +1,83 @@
+from typing import Dict, List
+from torch import LongTensor
+from transformers import DataCollatorForLanguageModeling
+
+
+class BERTIterator:
+
+    def __init__(self, dataset_reader, batch_size: int, sentence_len: int):
+        self.dataset_reader = dataset_reader
+        self.batch_size = batch_size
+        self.sentence_len = sentence_len
+
+        self.data_collator = DataCollatorForLanguageModeling(
+            tokenizer=self.dataset_reader.encoder.tokenizer_ref,
+            mlm = True,
+            mlm_probability = 0.15)
+
+    def load(self, dataset_meta) -> LongTensor:
+        self.dataset_reader.read(dataset_meta)
+
+        #In case user wants to display the data
+        return self.dataset_reader.encoded_text
+
+    def __iter__(self):
+
+        self.index = 0
+
+        return self
+
+    def __next__(self):
+
+        if self.index + self.batch_size > self.num_examples:
+            raise StopIteration
+
+        batch_examples = []
+
+        for i in range(self.batch_size):
+            example = self._load_example()
+            batch_examples.append(example)
+
+        batch = self._collate(batch_examples=batch_examples)
+
+        return batch
+
+    @property
+    def num_examples(self):
+        """Returns that number of non-overlapping  examples
+        in the dataset
+        """
+
+        num_examples = (len(self.dataset_reader.encoded_text) - 1) // self.sentence_len
+
+        return num_examples
+
+    @property
+    def num_batches(self):
+        """Returns the total number of batches. The last batch
+        is dropped if its size is less than self.batch_size.
+        """
+
+        num_batches = self.num_examples // self.batch_size
+
+        return num_batches
+
+    def _load_example(self) -> LongTensor:
+
+        # LongTensor containing the dataset
+        dataset = self.dataset_reader.encoded_text
+
+        #Getting an example - sequence of length 'sentence_len'
+        example = dataset.narrow(
+            dim=0, start=self.index * self.sentence_len, length=self.sentence_len
+        )
+
+        self.index += 1
+
+        return example
+
+    def _collate(self, batch_examples: List) -> Dict:
+
+        return self.data_collator(batch_examples)
+
+

syfertext/data/readers/language_modeling.py

@@ -21,7 +21,7 @@ def read(self, dataset_meta):
         data_path = pathlib.Path(data_path)
 
         # Open the text file to read and encode its text
-        with data_path.open() as f:
+        with data_path.open(encoding='utf-8') as f:
 
             # Read all lines
             for line in f.readlines():

syfertext/encoders/bert_encoder.py

@@ -0,0 +1,16 @@
+from typing import Dict, List
+from transformers import BertTokenizer
+
+class BERTEncoder:
+
+    def __init__(self):
+        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
+    def __call__(self, text:List) -> Dict:
+        inputs = self.tokenizer(text)
+        return {"token_ids": inputs["input_ids"]}
+
+    @property
+    def tokenizer_ref(self):
+        #decorator method so tokenizer can't be modified
+        return self.tokenizer