Files
multifit/tests/test_text_data.py
T

57 lines
2.0 KiB
Python

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)
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 == (lml.bs, 70)
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", 'xxbos', 'project', 'cool'])
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)
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 == (lml.bs, 70, 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], "xxbos fast ai is a cool project".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", 'xxbos', 'project', 'cool'])
###################### NEW CODE