| """Adapters for masked-MDLM and uniform-Duo Di4C students.""" |
|
|
| from __future__ import annotations |
|
|
| from pathlib import Path |
| import sys |
| from typing import Iterable |
|
|
| from dlb.adapters.base import AdapterError, BaseTeacherAdapter |
| from dlb.runner import RunRequest |
| from dlb.schema import SampleRecord |
|
|
|
|
| class Di4CAdapter(BaseTeacherAdapter): |
| identity = "dlb.adapters.di4c:v3" |
| upstream = "di4c" |
| teacher_families = { |
| "mdlm_di4c": "masked_mdlm", |
| "duo_di4c": "uniform_duo", |
| } |
| batch_sizes = { |
| ("mdlm_di4c", "lm1b"): 16, |
| ("mdlm_di4c", "owt"): 16, |
| ("duo_di4c", "lm1b"): 16, |
| ("duo_di4c", "owt"): 16, |
| } |
|
|
| def __init__(self, teacher: str | None = None) -> None: |
| if teacher not in {None, "mdlm", "duo"}: |
| raise AdapterError(f"unknown Di4C teacher identity: {teacher!r}") |
| self.supported_models = frozenset( |
| {"mdlm_di4c", "duo_di4c"} if teacher is None else {f"{teacher}_di4c"} |
| ) |
|
|
| def render_command( |
| self, request: RunRequest, run_dir: Path, *, dry_run: bool |
| ) -> list[str]: |
| root, length, batch_size = self._validate_request(request, run_dir) |
| checkpoint = self._resolve_checkpoint(root, request, dry_run=dry_run) |
| expected_family = self.teacher_families[request.model_id] |
| if checkpoint.teacher_family != expected_family: |
| raise AdapterError( |
| f"Di4C checkpoint teacher family {checkpoint.teacher_family!r} " |
| f"does not match {expected_family!r}" |
| ) |
| upstream_root = root / "upstreams" / "di4c" / "sdtt" |
| wrapper = root / "adapters" / "sample_di4c.py" |
| if not (upstream_root / "src/sdtt/main.py").is_file() or upstream_root.is_symlink(): |
| raise AdapterError(f"pinned Di4C language source is missing or unsafe: {upstream_root}") |
| if not wrapper.is_file() or wrapper.is_symlink(): |
| raise AdapterError(f"Di4C sampling wrapper is missing or unsafe: {wrapper}") |
| checkpoint_path = ( |
| checkpoint.path / "sdtt7-di4c2.ckpt" |
| if not checkpoint.path.suffix |
| else checkpoint.path |
| ) |
| config_path = checkpoint.config_path |
| if config_path is None: |
| raise AdapterError("Di4C checkpoint selection does not bind a sampling config") |
| checkpoint_sha256, config_sha256 = self._runtime_asset_digests( |
| root, |
| request, |
| checkpoint_path, |
| config_path, |
| dry_run=dry_run, |
| ) |
| tokenizer_snapshot = self._tokenizer_snapshot(root, request.dataset_id) |
| arguments = [ |
| sys.executable, |
| "-B", |
| "-u", |
| str(wrapper), |
| "--upstream-root", |
| str(upstream_root), |
| "--checkpoint", |
| str(checkpoint_path), |
| "--config", |
| str(config_path), |
| "--checkpoint-sha256", |
| checkpoint_sha256, |
| "--config-sha256", |
| config_sha256, |
| "--data-config", |
| str(root / "artifacts/data.yaml"), |
| "--downloads-manifest", |
| str(root / "data/manifests/downloads.json"), |
| "--tokenizer-snapshot", |
| str(tokenizer_snapshot), |
| "--output", |
| str(run_dir.resolve() / "upstream_token_ids.json"), |
| "--sample-count", |
| str(request.sample_count), |
| "--batch-size", |
| str(min(batch_size, request.sample_count)), |
| "--num-steps", |
| str(request.step_count), |
| "--seq-len", |
| str(length), |
| "--seed", |
| str(request.seed), |
| "--sampler", |
| "ancestral", |
| "--teacher-family", |
| checkpoint.teacher_family, |
| "--dataset", |
| request.dataset_id, |
| *self._conditional_script_flags(request), |
| "--offline", |
| "true", |
| ] |
| arguments.extend(["--allow-missing-embedded-config", "true"]) |
| self._validate_argv(arguments) |
| return arguments |
|
|
| def _tokenizer_snapshot(self, root: Path, dataset_id: str) -> Path: |
| tokenizer = self._load_data_contract(root, dataset_id)["tokenizer"] |
| revision = self._tokenizer_revision(root, tokenizer) |
| return ( |
| root |
| / "data/raw/huggingface/hub" |
| / f"models--{tokenizer}" |
| / "snapshots" |
| / revision |
| ) |
|
|
| def convert_outputs( |
| self, request: RunRequest, run_dir: Path |
| ) -> Iterable[SampleRecord]: |
| return self._convert_capture_outputs(request, run_dir) |
|
|
|
|
| adapter = Di4CAdapter() |
|
|
