| import json |
|
|
| import numpy as np |
|
|
| from nayara.data import MixedMemmap |
| from nayara.prepare_data import make_profile |
|
|
|
|
| def test_rank_disjoint_weighted_memmap(tmp_path): |
| sources = {} |
| for name, weight in (("code", 0.8), ("reasoning", 0.15), ("docs", 0.05)): |
| filename = f"{name}.bin" |
| np.arange(10_000, dtype=np.uint16).tofile(tmp_path / filename) |
| sources[name] = {"file": filename, "tokens": 10_000, "weight": weight} |
| (tmp_path / "manifest.json").write_text(json.dumps({"sources": sources}), encoding="utf-8") |
|
|
| rank0 = MixedMemmap(tmp_path, seq_len=16, batch_size=2, rank=0, world_size=2, seed=7) |
| rank1 = MixedMemmap(tmp_path, seq_len=16, batch_size=2, rank=1, world_size=2, seed=7) |
| x0, y0, source0 = rank0.next_cpu() |
| x1, y1, source1 = rank1.next_cpu() |
| assert source0 == source1 |
| assert x0.shape == y0.shape == (2, 16) |
| assert x1.shape == y1.shape == (2, 16) |
| assert not np.intersect1d(x0.numpy(), x1.numpy()).size |
|
|
|
|
| def test_profile_scales_mix_exactly(tmp_path): |
| source = tmp_path / "sources.json" |
| output = tmp_path / "scaled.json" |
| source.write_text(json.dumps({ |
| "code": {"weight": 0.80}, "reasoning": {"weight": 0.15}, "docs": {"weight": 0.05}, |
| }), encoding="utf-8") |
| make_profile(source, output, 101) |
| scaled = json.loads(output.read_text(encoding="utf-8")) |
| assert sum(item["target_tokens"] for item in scaled.values()) == 101 |
| assert scaled["code"]["target_tokens"] == 81 |
|
|
