194 lines
7.3 KiB
Python
194 lines
7.3 KiB
Python
"""Tests for pipeline.io module."""
|
||
|
||
import os
|
||
import tempfile
|
||
import pytest
|
||
import torch
|
||
import h5py
|
||
from pathlib import Path
|
||
|
||
from pipeline.io import FileScanner, HDF5Handler, cache_jsonl
|
||
|
||
|
||
class TestFileScanner:
|
||
def test_scan_files_in_directory(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
Path(tmpdir, "file1.txt").touch()
|
||
Path(tmpdir, "file2.txt").touch()
|
||
subdir = os.path.join(tmpdir, "subdir")
|
||
os.makedirs(subdir)
|
||
Path(subdir, "file3.txt").touch()
|
||
|
||
files = FileScanner.scan(tmpdir)
|
||
assert len(files) == 3
|
||
|
||
def test_scan_empty_directory(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
assert FileScanner.scan(tmpdir) == []
|
||
|
||
def test_scan_with_suffix_filter(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
Path(tmpdir, "file1.txt").touch()
|
||
Path(tmpdir, "file2.json").touch()
|
||
|
||
txt_files = FileScanner.scan(tmpdir, suffix=".txt")
|
||
assert len(txt_files) == 1
|
||
assert txt_files[0].endswith(".txt")
|
||
|
||
def test_scan_folders_in_directory(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
os.makedirs(os.path.join(tmpdir, "folder1"))
|
||
os.makedirs(os.path.join(tmpdir, "folder2"))
|
||
os.makedirs(os.path.join(tmpdir, "folder1", "nested"))
|
||
|
||
folders = FileScanner.scan_folders(tmpdir)
|
||
assert len(folders) == 3
|
||
|
||
def test_scan_folders_with_filter(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
os.makedirs(os.path.join(tmpdir, "folder1"))
|
||
os.makedirs(os.path.join(tmpdir, "folder2"))
|
||
|
||
folders = FileScanner.scan_folders(
|
||
tmpdir, filter_func=lambda x: "folder1" in x
|
||
)
|
||
assert len(folders) == 1
|
||
|
||
def test_group_by_extension(self):
|
||
files = ["/path/file1.txt", "/path/file2.txt", "/path/file3.json"]
|
||
groups = FileScanner.group_by_extension(files)
|
||
assert ".txt" in groups
|
||
assert ".json" in groups
|
||
assert len(groups[".txt"]) == 2
|
||
assert len(groups[".json"]) == 1
|
||
|
||
|
||
class TestHDF5Handler:
|
||
def test_save_and_load(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
tensor_group = {
|
||
"sequence": [torch.tensor([1, 2, 3], dtype=torch.int32)],
|
||
"labels": [torch.tensor([4, 5], dtype=torch.int32)],
|
||
}
|
||
h5_path = HDF5Handler.save(tmpdir, "test", tensor_group)
|
||
|
||
assert os.path.exists(h5_path)
|
||
loaded = HDF5Handler.load(h5_path, share_memory=False)
|
||
assert "sequence" in loaded
|
||
assert "labels" in loaded
|
||
assert torch.equal(
|
||
loaded["sequence"][0], torch.tensor([1, 2, 3], dtype=torch.int32)
|
||
)
|
||
assert torch.equal(
|
||
loaded["labels"][0], torch.tensor([4, 5], dtype=torch.int32)
|
||
)
|
||
|
||
def test_save_creates_directory(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
output_dir = os.path.join(tmpdir, "nested", "output")
|
||
HDF5Handler.save(output_dir, "test", {"data": [torch.tensor([1, 2, 3])]})
|
||
assert os.path.exists(output_dir)
|
||
assert os.path.exists(os.path.join(output_dir, "test.h5"))
|
||
|
||
def test_load_directory_with_multiple_files(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
for i, data in enumerate([[1, 2, 3], [4, 5, 6]]):
|
||
h5_path = os.path.join(tmpdir, f"file{i}.h5")
|
||
with h5py.File(h5_path, "w") as f:
|
||
grp = f.create_group("data")
|
||
grp.create_dataset("data_0", data=data)
|
||
|
||
loaded = HDF5Handler.load(tmpdir, share_memory=False)
|
||
assert len(loaded["data"]) == 2
|
||
|
||
def test_load_directory_with_nested_files(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
subdir = os.path.join(tmpdir, "subdir")
|
||
os.makedirs(subdir)
|
||
h5_path = os.path.join(subdir, "nested.h5")
|
||
with h5py.File(h5_path, "w") as f:
|
||
grp = f.create_group("test")
|
||
grp.create_dataset("data_0", data=[1, 2])
|
||
|
||
loaded = HDF5Handler.load(tmpdir, share_memory=False)
|
||
assert "test" in loaded
|
||
assert len(loaded["test"]) == 1
|
||
|
||
def test_save_multiple_tensors_per_key(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
tensor_group = {
|
||
"batch": [
|
||
torch.tensor([1, 2]),
|
||
torch.tensor([3, 4, 5]),
|
||
torch.tensor([6]),
|
||
],
|
||
}
|
||
HDF5Handler.save(tmpdir, "multi", tensor_group)
|
||
loaded = HDF5Handler.load(tmpdir, share_memory=False)
|
||
assert len(loaded["batch"]) == 3
|
||
|
||
def test_get_metadata(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
tensor_group = {
|
||
"data": [torch.tensor([1, 2, 3]) for _ in range(5)],
|
||
}
|
||
HDF5Handler.save(tmpdir, "meta", tensor_group)
|
||
|
||
h5_path = os.path.join(tmpdir, "meta.h5")
|
||
metadata = HDF5Handler.get_metadata(h5_path)
|
||
assert metadata["data"] == 5
|
||
|
||
|
||
class DummyTokenizer:
|
||
def __init__(self):
|
||
self._special_token_map = {}
|
||
self._chat_template = None
|
||
|
||
def encode(self, text: str, add_special_tokens: bool = False):
|
||
return [ord(c) for c in text]
|
||
|
||
def decode(self, tokens, skip_special_tokens=True):
|
||
return "".join(chr(t) for t in tokens)
|
||
|
||
def token_to_id(self, token: str):
|
||
return ord(token)
|
||
|
||
def set_chat_template(self, template):
|
||
self._chat_template = template
|
||
|
||
def apply_chat_template(self, messages, add_generation_prompt=True, tokenize=True):
|
||
text = ""
|
||
for m in messages:
|
||
text += f"<｜im▁start｜>{m['role']}\n{m['content']}<｜im▁end｜>\n"
|
||
if add_generation_prompt:
|
||
text += "<｜im▁start｜>assistant\n"
|
||
return self.encode(text) if tokenize else text
|
||
|
||
|
||
class TestPositionIds:
|
||
def test_example_specific_position_ids(self):
|
||
with tempfile.TemporaryDirectory() as tmpdir:
|
||
jsonl_path = os.path.join(tmpdir, "data.jsonl")
|
||
with open(jsonl_path, "w") as f:
|
||
f.write('{"messages": [{"role": "user", "content": "a"}, {"role": "assistant", "content": "bc"}]}\n')
|
||
f.write('{"messages": [{"role": "user", "content": "def"}, {"role": "assistant", "content": "g"}]}\n')
|
||
|
||
from pipeline.processors import SFTProcessor
|
||
|
||
processor = SFTProcessor(DummyTokenizer())
|
||
out_dir = os.path.join(tmpdir, "cached")
|
||
cache_jsonl([jsonl_path], out_dir, processor, pack_size=-1, group_size=0)
|
||
|
||
h5_path = os.path.join(out_dir, "data.h5")
|
||
loaded = HDF5Handler.load(h5_path, share_memory=False)
|
||
|
||
assert "position_ids" in loaded
|
||
assert len(loaded["position_ids"]) == 2
|
||
assert len(loaded["position_ids"]) == len(loaded["sequence"])
|
||
|
||
for i, pos in enumerate(loaded["position_ids"]):
|
||
seq_len = len(loaded["sequence"][i])
|
||
assert len(pos) == seq_len
|
||
assert pos[0].item() == 0
|
||
assert (pos == torch.arange(seq_len, dtype=torch.int32)).all()
|