ViperEkura/DataPipeline
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 修复特殊token 问题

Browse Source
This commit is contained in:
ViperEkura
2026-04-02 16:16:02 +08:00
parent e01ec081b3
commit f44ad6912e
28 changed files with 334 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/test_io.py
+18 -9
View File
@@ -11,7 +11,6 @@ from pipeline.io import IOHandler
class TestIOHandler:
def test_fetch_files_in_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
Path(tmpdir, "file1.txt").touch()
@@ -41,7 +40,9 @@ class TestIOHandler:
os.makedirs(os.path.join(tmpdir, "folder1"))
os.makedirs(os.path.join(tmpdir, "folder2"))
folders = IOHandler.fetch_folders(tmpdir, filter_func=lambda x: "folder1" in x)
folders = IOHandler.fetch_folders(
tmpdir, filter_func=lambda x: "folder1" in x
)
assert len(folders) == 1
def test_save_and_load_h5(self):
@@ -56,8 +57,12 @@ class TestIOHandler:
loaded = IOHandler.load_h5(tmpdir, 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))
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_h5_creates_directory(self):
with tempfile.TemporaryDirectory() as tmpdir:
@@ -70,9 +75,9 @@ class TestIOHandler:
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:
with h5py.File(h5_path, "w") as f:
grp = f.create_group("data")
grp.create_dataset('data_0', data=data)
grp.create_dataset("data_0", data=data)
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
assert len(loaded["data"]) == 2
@@ -82,9 +87,9 @@ class TestIOHandler:
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:
with h5py.File(h5_path, "w") as f:
grp = f.create_group("test")
grp.create_dataset('data_0', data=[1, 2])
grp.create_dataset("data_0", data=[1, 2])
loaded = IOHandler.load_h5(tmpdir, share_memory=False)
assert "test" in loaded
@@ -93,7 +98,11 @@ class TestIOHandler:
def test_save_h5_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])],
"batch": [
torch.tensor([1, 2]),
torch.tensor([3, 4, 5]),
torch.tensor([6]),
],
}
IOHandler.save_h5(tmpdir, "multi", tensor_group)
loaded = IOHandler.load_h5(tmpdir, share_memory=False)