File size: 3,817 Bytes
79a20c7 | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | 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
|