| from __future__ import annotations |
|
|
| import json |
| import hashlib |
| from pathlib import Path |
|
|
| import pytest |
| import torch |
| from safetensors.torch import save_file |
|
|
| from offload.dspark.d2.train.data import D2Blocks, sha256 |
|
|
|
|
| def write_shard(root: Path, gen: int = 4, split: str = "train", L: int = 10) -> Path: |
| shard = root/(split+".safetensors") |
| p = torch.full((L-gen, 20), .045) |
| save_file({"seq_index": torch.tensor([[0, L, gen, 0, 0, L, 0, L-gen]]), |
| "tokens": torch.arange(L, dtype=torch.int32), "taps": torch.zeros(L, 15360, dtype=torch.bfloat16), |
| "row_pos": torch.arange(L, dtype=torch.int32), |
| "top_ids": torch.arange(20, dtype=torch.int32).expand(L-gen, -1).contiguous(), "top_lp": p.log(), |
| "resid": 1-p.sum(-1), "tok_lp": p[:, 0].log(), "tok_rank": torch.ones(L-gen, dtype=torch.int32)}, |
| shard, metadata={"schema": "d2.shard.v1", "split": split}) |
| meta = {"schema": "d2.shard.v1", "split": split, "W": 128, "K": 5, "topk": 20, |
| "tap_layers": [37, 38, 39], "hidden": 5120, "noise_token": 128799, "sha256": sha256(shard), |
| "sequences": [{"seq_id": "p1", "source_docs": ["doc1"], "family": "argument_section", "L": L, "gen_start": gen, |
| "tokens_sha256": hashlib.sha256(torch.arange(L, dtype=torch.int32).numpy().tobytes()).hexdigest()}]} |
| side = root/(split+".json") |
| side.write_text(json.dumps(meta)) |
| manifest = root/"MANIFEST.json" |
| manifest.write_text(json.dumps({"schema": "d2.manifest.v1", "checks": {"plumbing_errors": 0}, |
| "splits": {split: {"files": [{"file": shard.name, "sha256": sha256(shard), "sidecar": side.name, |
| "sidecar_sha256": sha256(side)}]}}})) |
| return manifest |
|
|
|
|
| def test_normative_target_offset_and_partial_tail(tmp_path: Path) -> None: |
| d = D2Blocks(write_shard(tmp_path), "train") |
| assert len(d) == 5 |
| first, last = d[0], d[len(d)-1] |
| assert first["inputs"]["P"].item() == 3 |
| assert first["inputs"]["anchor"].item() == 4 |
| assert first["targets"]["labels"].tolist() == [[5, 6, 7, 8, 9]] |
| assert first["targets"]["previous"].tolist() == [[4, 5, 6, 7, 8]] |
| assert last["inputs"]["P"].item() == 7 |
| assert last["targets"]["labels"].tolist() == [[9, -1, -1, -1, -1]] |
| assert last["targets"]["position_mask"].tolist() == [[True, False, False, False, False]] |
| assert last["inputs"]["positions"].tolist() == [list(range(8))] |
|
|
|
|
| def test_changed_shard_refused(tmp_path: Path) -> None: |
| p = write_shard(tmp_path) |
| (tmp_path/"train.json").write_text("{}") |
| with pytest.raises(ValueError, match="hash/path mismatch"): |
| D2Blocks(p, "train") |
|
|
