File size: 1,472 Bytes
0990c3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | 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
|