| """Unit tests for dataset loaders.""" | |
| import pytest | |
| from aegis.data.loaders import get_data_loader, GSM8KLoader, ToxiGenLoader, HC3Loader | |
| from aegis.data.loaders.base import DataSample | |
| def test_gsm8k_loader(): | |
| loader = GSM8KLoader(split="train", max_samples=5) | |
| samples = loader.get_samples() | |
| assert len(samples) == 5 | |
| for s in samples: | |
| assert isinstance(s, DataSample) | |
| assert len(s.input) > 0 | |
| assert len(s.label) > 0 | |
| assert "dataset" in s.metadata | |
| def test_toxigen_loader(): | |
| loader = ToxiGenLoader(split="train", max_samples=4) | |
| samples = loader.get_samples() | |
| assert len(samples) == 4 | |
| for s in samples: | |
| assert isinstance(s, DataSample) | |
| assert s.label in ["toxic", "non-toxic"] | |
| def test_hc3_loader(): | |
| loader = HC3Loader(split="train", max_samples=4) | |
| samples = loader.get_samples() | |
| assert len(samples) == 4 | |
| for s in samples: | |
| assert isinstance(s, DataSample) | |
| assert s.label in ["human", "ai"] | |
| def test_get_data_loader_factory(): | |
| loader_gsm = get_data_loader("gsm8k", split="test", max_samples=2) | |
| assert isinstance(loader_gsm, GSM8KLoader) | |
| assert len(loader_gsm) == 2 | |
| loader_tox = get_data_loader("toxigen", split="test", max_samples=2) | |
| assert isinstance(loader_tox, ToxiGenLoader) | |
| assert len(loader_tox) == 2 | |
| loader_hc3 = get_data_loader("hc3", split="test", max_samples=2) | |
| assert isinstance(loader_hc3, HC3Loader) | |
| assert len(loader_hc3) == 2 | |
Xet Storage Details
- Size:
- 1.5 kB
- Xet hash:
- 75f41ffc73dc39da447b1f3fbae0782e277c25ded7519e6a094f46cefcde64e7
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.