def test_as_dataset(self):
with tempfile.TemporaryDirectory() as tmp_dir:
dummy_builder = DummyBuilder(cache_dir=tmp_dir, name="dummy")
os.makedirs(dummy_builder.cache_dir)
dummy_builder.info.splits = SplitDict()
dummy_builder.info.splits.add(SplitInfo("train", num_examples=10))
dummy_builder.info.splits.add(SplitInfo("test", num_examples=10))
for split in dummy_builder.info.splits:
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"dummy_builder-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 10})
writer.finalize()
dsets = dummy_builder.as_dataset()
self.assertIsInstance(dsets, DatasetDict)
self.assertListEqual(list(dsets.keys()), ["train", "test"])
self.assertEqual(len(dsets["train"]), 10)
self.assertEqual(len(dsets["test"]), 10)
dset = dummy_builder.as_dataset("train")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train")
self.assertEqual(len(dset), 10)
dset = dummy_builder.as_dataset("train+test[:30%]")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train+test[:30%]")
self.assertEqual(len(dset), 13)
def test_read(self):
name = "my_name"
train_info = SplitInfo(name="train", num_examples=100)
test_info = SplitInfo(name="test", num_examples=100)
split_infos = [train_info, test_info]
split_dict = SplitDict()
split_dict.add(train_info)
split_dict.add(test_info)
info = DatasetInfo(splits=split_dict)
reader = ReaderTest("", info)
instructions = "test[:33%]"
dset = reader.read(name, instructions, split_infos)
self.assertEqual(dset["filename"][0], f"{name}-test")
self.assertEqual(dset.num_rows, 33)
self.assertEqual(dset.num_columns, 1)
instructions = ["train", "test[:33%]"]
train_dset, test_dset = reader.read(name, instructions, split_infos)
self.assertEqual(train_dset["filename"][0], f"{name}-train")
self.assertEqual(train_dset.num_rows, 100)
self.assertEqual(train_dset.num_columns, 1)
self.assertEqual(test_dset["filename"][0], f"{name}-test")
self.assertEqual(test_dset.num_rows, 33)
self.assertEqual(test_dset.num_columns, 1)
def _create_dummy_dataset(self):
name = "my_name"
train_info = SplitInfo(name="train", num_examples=30)
test_info = SplitInfo(name="test", num_examples=30)
split_infos = [train_info, test_info]
split_dict = SplitDict()
split_dict.add(train_info)
split_dict.add(test_info)
info = DatasetInfo(splits=split_dict)
reader = ReaderTester("", info)
dset = reader.read(name, "train", split_infos)
return dset
def test_read_files(self):
train_info = SplitInfo(name="train", num_examples=100)
test_info = SplitInfo(name="test", num_examples=100)
split_dict = SplitDict()
split_dict.add(train_info)
split_dict.add(test_info)
info = DatasetInfo(splits=split_dict)
reader = ReaderTest("", info)
files = [{"filename": "train"}, {"filename": "test", "skip": 10, "take": 10}]
dset = reader.read_files(files, original_instructions="")
self.assertEqual(dset.num_rows, 110)
self.assertEqual(dset.num_columns, 1)
self.assertEqual(dset._data_files, files)
def test_as_dataset_with_post_process(self):
def _post_process(self, dataset, resources_paths):
def char_tokenize(example):
return {"tokens": list(example["text"])}
return dataset.map(
char_tokenize,
cache_file_name=resources_paths["tokenized_dataset"])
def _post_processing_resources(self, split):
return {
"tokenized_dataset":
"tokenized_dataset-{split}.arrow".format(split=split)
}
with tempfile.TemporaryDirectory() as tmp_dir:
dummy_builder = DummyBuilder(cache_dir=tmp_dir, name="dummy")
dummy_builder.info.post_processed = PostProcessedInfo(
features=Features({
"text": Value("string"),
"tokens": [Value("string")]
}))
dummy_builder._post_process = types.MethodType(
_post_process, dummy_builder)
dummy_builder._post_processing_resources = types.MethodType(
_post_processing_resources, dummy_builder)
os.makedirs(dummy_builder.cache_dir)
dummy_builder.info.splits = SplitDict()
dummy_builder.info.splits.add(SplitInfo("train", num_examples=10))
dummy_builder.info.splits.add(SplitInfo("test", num_examples=10))
for split in dummy_builder.info.splits:
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"dummy_builder-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 10})
writer.finalize()
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"tokenized_dataset-{split}.arrow"),
features=Features({
"text": Value("string"),
"tokens": [Value("string")]
}),
)
writer.write_batch({
"text": ["foo"] * 10,
"tokens": [list("foo")] * 10
})
writer.finalize()
dsets = dummy_builder.as_dataset()
self.assertIsInstance(dsets, DatasetDict)
self.assertListEqual(list(dsets.keys()), ["train", "test"])
self.assertEqual(len(dsets["train"]), 10)
self.assertEqual(len(dsets["test"]), 10)
self.assertDictEqual(
dsets["train"].features,
Features({
"text": Value("string"),
"tokens": [Value("string")]
}))
self.assertDictEqual(
dsets["test"].features,
Features({
"text": Value("string"),
"tokens": [Value("string")]
}))
self.assertListEqual(dsets["train"].column_names,
["text", "tokens"])
self.assertListEqual(dsets["test"].column_names,
["text", "tokens"])
dset = dummy_builder.as_dataset("train")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train")
self.assertEqual(len(dset), 10)
self.assertDictEqual(
dset.features,
Features({
"text": Value("string"),
"tokens": [Value("string")]
}))
self.assertListEqual(dset.column_names, ["text", "tokens"])
self.assertGreater(dummy_builder.info.post_processing_size, 0)
self.assertGreater(
dummy_builder.info.post_processed.resources_checksums["train"]
["tokenized_dataset"]["num_bytes"], 0)
dset = dummy_builder.as_dataset("train+test[:30%]")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train+test[:30%]")
self.assertEqual(len(dset), 13)
self.assertDictEqual(
dset.features,
Features({
"text": Value("string"),
"tokens": [Value("string")]
}))
self.assertListEqual(dset.column_names, ["text", "tokens"])
def _post_process(self, dataset, resources_paths):
return dataset.select([0, 1], keep_in_memory=True)
with tempfile.TemporaryDirectory() as tmp_dir:
dummy_builder = DummyBuilder(cache_dir=tmp_dir, name="dummy")
dummy_builder._post_process = types.MethodType(
_post_process, dummy_builder)
os.makedirs(dummy_builder.cache_dir)
dummy_builder.info.splits = SplitDict()
dummy_builder.info.splits.add(SplitInfo("train", num_examples=10))
dummy_builder.info.splits.add(SplitInfo("test", num_examples=10))
for split in dummy_builder.info.splits:
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"dummy_builder-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 10})
writer.finalize()
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"small_dataset-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 2})
writer.finalize()
dsets = dummy_builder.as_dataset()
self.assertIsInstance(dsets, DatasetDict)
self.assertListEqual(list(dsets.keys()), ["train", "test"])
self.assertEqual(len(dsets["train"]), 2)
self.assertEqual(len(dsets["test"]), 2)
self.assertDictEqual(dsets["train"].features,
Features({"text": Value("string")}))
self.assertDictEqual(dsets["test"].features,
Features({"text": Value("string")}))
self.assertListEqual(dsets["train"].column_names, ["text"])
self.assertListEqual(dsets["test"].column_names, ["text"])
dset = dummy_builder.as_dataset("train")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train")
self.assertEqual(len(dset), 2)
self.assertDictEqual(dset.features,
Features({"text": Value("string")}))
self.assertListEqual(dset.column_names, ["text"])
dset = dummy_builder.as_dataset("train+test[:30%]")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train+test[:30%]")
self.assertEqual(len(dset), 2)
self.assertDictEqual(dset.features,
Features({"text": Value("string")}))
self.assertListEqual(dset.column_names, ["text"])
def _post_process(self, dataset, resources_paths):
if os.path.exists(resources_paths["index"]):
dataset.load_faiss_index("my_index", resources_paths["index"])
return dataset
else:
dataset.add_faiss_index_from_external_arrays(
external_arrays=np.ones((len(dataset), 8)),
string_factory="Flat",
index_name="my_index")
dataset.save_faiss_index("my_index", resources_paths["index"])
return dataset
def _post_processing_resources(self, split):
return {"index": "Flat-{split}.faiss".format(split=split)}
with tempfile.TemporaryDirectory() as tmp_dir:
dummy_builder = DummyBuilder(cache_dir=tmp_dir, name="dummy")
dummy_builder._post_process = types.MethodType(
_post_process, dummy_builder)
dummy_builder._post_processing_resources = types.MethodType(
_post_processing_resources, dummy_builder)
os.makedirs(dummy_builder.cache_dir)
dummy_builder.info.splits = SplitDict()
dummy_builder.info.splits.add(SplitInfo("train", num_examples=10))
dummy_builder.info.splits.add(SplitInfo("test", num_examples=10))
for split in dummy_builder.info.splits:
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"dummy_builder-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 10})
writer.finalize()
writer = ArrowWriter(
path=os.path.join(dummy_builder.cache_dir,
f"small_dataset-{split}.arrow"),
features=Features({"text": Value("string")}),
)
writer.write_batch({"text": ["foo"] * 2})
writer.finalize()
dsets = dummy_builder.as_dataset()
self.assertIsInstance(dsets, DatasetDict)
self.assertListEqual(list(dsets.keys()), ["train", "test"])
self.assertEqual(len(dsets["train"]), 10)
self.assertEqual(len(dsets["test"]), 10)
self.assertDictEqual(dsets["train"].features,
Features({"text": Value("string")}))
self.assertDictEqual(dsets["test"].features,
Features({"text": Value("string")}))
self.assertListEqual(dsets["train"].column_names, ["text"])
self.assertListEqual(dsets["test"].column_names, ["text"])
self.assertListEqual(dsets["train"].list_indexes(), ["my_index"])
self.assertListEqual(dsets["test"].list_indexes(), ["my_index"])
self.assertGreater(dummy_builder.info.post_processing_size, 0)
self.assertGreater(
dummy_builder.info.post_processed.resources_checksums["train"]
["index"]["num_bytes"], 0)
dset = dummy_builder.as_dataset("train")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train")
self.assertEqual(len(dset), 10)
self.assertDictEqual(dset.features,
Features({"text": Value("string")}))
self.assertListEqual(dset.column_names, ["text"])
self.assertListEqual(dset.list_indexes(), ["my_index"])
dset = dummy_builder.as_dataset("train+test[:30%]")
self.assertIsInstance(dset, Dataset)
self.assertEqual(dset.split, "train+test[:30%]")
self.assertEqual(len(dset), 13)
self.assertDictEqual(dset.features,
Features({"text": Value("string")}))
self.assertListEqual(dset.column_names, ["text"])
self.assertListEqual(dset.list_indexes(), ["my_index"])