| import json |
| from pathlib import Path |
| import sys |
| from types import SimpleNamespace |
| from types import ModuleType |
|
|
| import dlb.adapters.capture as capture_module |
|
|
|
|
| def test_load_entrypoint_registers_module_for_hydra_relative_configs( |
| tmp_path: Path, |
| ) -> None: |
| """Catch Hydra entrypoints losing module metadata for relative config paths.""" |
|
|
| entrypoint = tmp_path / "main.py" |
| entrypoint.write_text( |
| """ |
| import sys |
| |
| MODULE_FILE_AT_IMPORT = getattr(sys.modules.get(__name__), "__file__", None) |
| |
| def main(): |
| pass |
| """, |
| encoding="utf-8", |
| ) |
|
|
| try: |
| module = capture_module._load_entrypoint(entrypoint) |
| finally: |
| sys.modules.pop("dlb_pinned_upstream_main", None) |
|
|
| assert Path(module.MODULE_FILE_AT_IMPORT) == entrypoint |
|
|
|
|
| def test_run_main_points_hydra_at_sibling_config_directory(tmp_path: Path) -> None: |
| """Catch dynamic imports making Hydra treat sibling configs as a missing package.""" |
|
|
| entrypoint = tmp_path / "main.py" |
| config_dir = tmp_path / "configs" |
| config_dir.mkdir() |
| entrypoint.write_text("# fixture\n", encoding="utf-8") |
| observed: dict[str, list[str]] = {} |
| module = ModuleType("fixture_upstream") |
|
|
| def main() -> None: |
| observed["argv"] = list(sys.argv) |
|
|
| module.main = main |
|
|
| capture_module._run_main(module, entrypoint, ["mode=sample_eval"]) |
|
|
| assert observed["argv"] == [ |
| str(entrypoint), |
| f"--config-path={config_dir.resolve()}", |
| "mode=sample_eval", |
| ] |
|
|
|
|
| def test_teacher_capture_accepts_upstream_python_list_samples(tmp_path: Path) -> None: |
| """Catch FLM restore_model_and_sample returning list rows instead of a tensor.""" |
|
|
| class Tokenizer: |
| def batch_decode(self, rows): |
| assert rows == [[1, 2], [3, 4]] |
| return ["one", "two"] |
|
|
| class TrainerBase: |
| tokenizer = Tokenizer() |
|
|
| def restore_model_and_sample(self): |
| return [[1, 2], [3, 4]] |
|
|
| module = ModuleType("fixture_teacher") |
| module.algo = SimpleNamespace( |
| trainer_base=SimpleNamespace(TrainerBase=TrainerBase) |
| ) |
| module.main = lambda: TrainerBase().restore_model_and_sample() |
| capture_path = tmp_path / "capture.json" |
| invocation = capture_module.CaptureInvocation( |
| entrypoint=tmp_path / "main.py", |
| capture_path=capture_path, |
| ) |
|
|
| capture_module._capture_teacher(module, invocation, []) |
|
|
| assert json.loads(capture_path.read_text(encoding="utf-8")) == { |
| "samples": [ |
| {"sample_id": 0, "text": "one", "token_ids": [1, 2]}, |
| {"sample_id": 1, "text": "two", "token_ids": [3, 4]}, |
| ], |
| "schema": "dlb-upstream-token-capture-v1", |
| } |
|
|
|
|
| def test_hf_masked_lm_capture_adapter_translates_duo_backbone_keywords( |
| monkeypatch, |
| ) -> None: |
| """Catch HF OWT checkpoints receiving the upstream x/sigma call contract.""" |
|
|
| calls = [] |
|
|
| class Output: |
| logits = "logits" |
|
|
| class HuggingFaceBackbone: |
| def forward(self, input_ids, timesteps): |
| calls.append((input_ids, timesteps)) |
| return Output() |
|
|
| class AutoModelForMaskedLM: |
| @staticmethod |
| def from_pretrained(path): |
| return HuggingFaceBackbone() |
|
|
| fake_transformers = SimpleNamespace(AutoModelForMaskedLM=AutoModelForMaskedLM) |
| monkeypatch.setitem(sys.modules, "transformers", fake_transformers) |
|
|
| with capture_module._patched_hf_masked_lm_backbone(): |
| backbone = fake_transformers.AutoModelForMaskedLM.from_pretrained("checkpoint") |
| output = backbone(x="tokens", sigma="times", class_cond=None, weights=None) |
|
|
| assert output == "logits" |
| assert calls == [("tokens", "times")] |
| restored = fake_transformers.AutoModelForMaskedLM.from_pretrained("checkpoint") |
| assert type(restored) is HuggingFaceBackbone |
|
|
