diff --git a/fastai_contrib/data.py b/fastai_contrib/data.py new file mode 100644 index 0000000..74c6b6c --- /dev/null +++ b/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 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 diff --git a/tests/test_text_data.py b/tests/test_text_data.py new file mode 100644 index 0000000..6b3dabf --- /dev/null +++ b/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 +