Add BiLM LanguageModelLoader with tests

fastai_contrib/data.py

@@ -0,0 +1,60 @@
+"NLP data loading pipeline. Supports csv, folders, and preprocessed data."
+from fastai.text import *
+from fastai.torch_core import *
+from fastai.text.transform import *
+from fastai.basic_data import *
+from fastai.data_block import *
+
+
+###################### UPDATED CODE
+
+LanguageModelType=Enum('LanguageModelType', 'FwdLM BwdLM BiLM')
+
+class LanguageModelLoader(): # copy of the original LanguageModelLoader
+    "Create a dataloader with bptt slightly changing."
+    def __init__(self, dataset:LabelList, bs:int=64, bptt:int=70,
+                 lm_type:LanguageModelType=LanguageModelType.FwdLM, shuffle:bool=False,
+                 max_len:int=25):
+        self.dataset,self.bs,self.bptt,self.lm_type,self.shuffle = dataset,bs,bptt,lm_type,shuffle
+        self.first,self.i,self.iter = True,0,0
+        self.n = len(np.concatenate(dataset.x.items)) // self.bs
+        self.max_len,self.num_workers = max_len,0
+
+    def __iter__(self):
+        if getattr(self.dataset, 'item', None) is not None:
+            yield LongTensor(getattr(self.dataset, 'item')).unsqueeze(1),LongTensor([0])
+        idx = np.random.permutation(len(self.dataset)) if self.shuffle else range(len(self.dataset))
+        data = self.batchify(np.concatenate([self.dataset.x[i] for i in idx]))
+
+        pos, itr = 0,0
+        while pos < self.n-1 and itr<len(self):
+            if self.first and pos == 0: self.first,seq_len = False,self.bptt + self.max_len
+            else:
+                bptt = self.bptt if np.random.random() < 0.95 else self.bptt / 2.
+                seq_len = max(5, int(np.random.normal(bptt, 5)))
+                seq_len = min(seq_len, self.bptt + self.max_len)
+            res = self.get_batch(data, pos, seq_len)
+            pos += seq_len
+            itr += 1
+            yield res
+
+    def __len__(self) -> int: return (self.n-1) // self.bptt
+    def __getattr__(self,k:str)->Any: return getattr(self.dataset, k)
+
+    def batchify(self, data:np.ndarray) -> LongTensor:
+        "Split the corpus `data` in batches."
+        nb = data.shape[0] // self.bs
+        data = np.array(data[:nb*self.bs]).reshape(self.bs, -1).T
+        if self.lm_type == LanguageModelType.BwdLM: data=data[::-1].copy()
+        elif self.lm_type == LanguageModelType.BiLM: data = np.stack([data, data[::-1].copy()], axis=2)
+        return LongTensor(data)
+
+    def get_batch(self, data:LongTensor,  i:int, seq_len:int) -> Tuple[LongTensor, LongTensor]:
+        "Create a batch at `i` of a given `seq_len`."
+        seq_len = min(seq_len, len(self.data) - 1 - i)
+        x = data[i:i+seq_len]
+        y = data[i+1:i+1+seq_len].contiguous() # x & y has 2 elements on the last dimension
+        y = y.view(-1, 2) if self.lm_type == LanguageModelType.BiLM else y.view(-1)
+        return x,y
+
+###################### NEW CODE

tests/test_text_data.py

@@ -0,0 +1,59 @@
+import pytest
+import fastai.text
+
+from fastai import *
+from fastai.text import *
+
+import fastai_contrib.data as contrib_data
+
+def text_df(labels):
+    data = []
+    texts = ["fast ai is a cool project", "hello world"] * 20
+    for ind, text in enumerate(texts):
+        sample = {}
+        sample["label"] = labels[ind%len(labels)]
+        sample["text"] = text
+        data.append(sample)
+    return pd.DataFrame(data)
+
+###################### UPDATED CODE
+def test_should_load_backwards_lm():
+    path = untar_data(URLs.IMDB_SAMPLE)
+    df = text_df(['neg','pos'])
+
+    data = TextLMDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=0, text_cols=["text"], bs=2,
+                                   lm_type=contrib_data.LanguageModelType.BwdLM,
+                                   ld_cls=contrib_data.LanguageModelLoader)
+    lml = data.train_dl.dl
+    lml.data = lml.batchify(np.concatenate([lml.dataset.x.items[i] for i in range(len(lml.dataset))]))
+    batch = lml.get_batch(lml.data, 0, 70)
+
+    assert batch[0].shape == (70, lml.bs)
+    assert batch[1].shape == (70*lml.bs,)
+
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:,0]]
+    np.testing.assert_array_equal(as_text[:5], ["world", "hello", '1', 'xxfld', 'project',])
+
+def test_should_load_bi_lm():
+    path = untar_data(URLs.IMDB_SAMPLE)
+    df = text_df(['neg', 'pos'])
+
+    data = TextLMDataBunch.from_df(path, train_df=df, valid_df=df, label_cols=0, text_cols=["text"], bs=2,
+                                   lm_type=contrib_data.LanguageModelType.BiLM,
+                                   ld_cls=contrib_data.LanguageModelLoader)
+    lml = data.train_dl.dl
+    lml.data = lml.batchify(np.concatenate([lml.dataset.x.items[i] for i in range(len(lml.dataset))]))
+    batch = lml.get_batch(lml.data, 0, 70)
+
+    assert batch[0].shape == (70, lml.bs, 2)
+    assert batch[1].shape == (70*lml.bs, 2)
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:, 0, 0]]
+    np.testing.assert_array_equal(as_text[:7], "xxfld 1 fast ai is a cool".split())
+
+    as_text = [lml.dataset.vocab.itos[x] for x in batch[0][:,0,1]]
+    np.testing.assert_array_equal(as_text[:5], ["world", "hello", '1', 'xxfld', 'project',])
+
+###################### NEW CODE
+