Blame for src/dlb/adapters/capture.py

"""Run a pinned upstream entrypoint while capturing returned token-ID samples."""

from __future__ import annotations

import builtins
from contextlib import contextmanager
from contextlib import ExitStack
from contextlib import nullcontext
from dataclasses import dataclass
import importlib
import importlib.util
import inspect
import json
import math
import os
from pathlib import Path
import sys
from types import ModuleType
from typing import Callable

import yaml

from dlb.io import atomic_json_write
from dlb.adapters.conditional_runtime import (
    adapt_candi_generate_sample_prompt,
    candi_prompt_conditioning,
    clamp_token_prefix,
    embedding_project_fn,
    load_conditioning_batch,
    patched_attribute,
    rdlm_project_fn,
    token_project_fn,
    vocab_project_fn,
)
from dlb.timing import benchmark_and_publish


@dataclass(frozen=True)
class CaptureInvocation:
    entrypoint: Path
    capture_path: Path
    kind: str = "teacher"
    expected_samples: int | None = None
    saved_config_path: Path | None = None
    saved_sde_path: Path | None = None
    data_config_path: Path | None = None
    downloads_manifest_path: Path | None = None
    dataset_id: str | None = None
    tokenizer_snapshot: Path | None = None
    benchmark_output: Path | None = None
    benchmark_metadata: Path | None = None
    benchmark_precision: str | None = None
    generation_mode: str = "unconditional"
    conditioning_manifest: Path | None = None
    conditioning_manifest_sha256: str | None = None
    conditioning_config_sha256: str | None = None
    prefix_length: int | None = None
    evaluation_continuation_length: int | None = None
    prompt_count: int | None = None
    diversity_prompt_count: int | None = None
    completions_per_diversity_prompt: int | None = None
    completion_schedule: str | None = None


@dataclass(frozen=True)
class TokenizerBinding:
    tokenizer_id: str
    revision: str
    snapshot: Path


def _split_arguments(argv: list[str]) -> tuple[CaptureInvocation, list[str]]:
    try:
        separator = argv.index("--")
    except ValueError as error:
        raise ValueError("capture wrapper requires -- before upstream arguments") from error
    wrapper = argv[:separator]
    forwarded = argv[separator + 1 :]
    allowed = {
        "--upstream-entrypoint",
        "--capture-path",
        "--capture-kind",
        "--expected-samples",
        "--saved-config-path",
        "--saved-sde-path",
        "--data-config-path",
        "--downloads-manifest-path",
        "--dataset-id",
        "--tokenizer-snapshot",
        "--benchmark-output",
        "--benchmark-metadata",
        "--benchmark-precision",
        "--generation-mode",
        "--conditioning-manifest",
        "--conditioning-manifest-sha256",
        "--conditioning-config-sha256",
        "--prefix-length",
        "--evaluation-continuation-length",
        "--prompt-count",
        "--diversity-prompt-count",
        "--completions-per-diversity-prompt",
        "--completion-schedule",
    }
    values: dict[str, str] = {}
    for argument in wrapper:
        key, marker, value = argument.partition("=")
        if not marker or key not in allowed or not value:
            raise ValueError(f"invalid capture-wrapper argument: {argument!r}")
        if key in values:
            raise ValueError(f"duplicate capture-wrapper argument: {key}")
        values[key] = value
    required = {"--upstream-entrypoint", "--capture-path"}
    if not required.issubset(values):
        raise ValueError("capture wrapper requires entrypoint and capture path")

    entrypoint = Path(values["--upstream-entrypoint"])
    capture_path = Path(values["--capture-path"])
    if not entrypoint.is_absolute() or entrypoint.is_symlink() or not entrypoint.is_file():
        raise ValueError(f"upstream entrypoint is missing or unsafe: {entrypoint}")
    if not capture_path.is_absolute():
        raise ValueError("capture path must be absolute")
    kind = values.get("--capture-kind", "teacher")
    if kind not in {"teacher", "langflow", "rdlm"}:
        raise ValueError(f"unsupported capture kind: {kind!r}")
    expected_text = values.get("--expected-samples")
    if expected_text is None:
        expected_samples = None
    else:
        try:
            expected_samples = int(expected_text)
        except ValueError as error:
            raise ValueError("expected sample count must be an integer") from error
        if expected_samples <= 0:
            raise ValueError("expected sample count must be positive")
    saved_config_path = (
        Path(values["--saved-config-path"])
        if "--saved-config-path" in values
        else None
    )
    saved_sde_path = (
        Path(values["--saved-sde-path"])
        if "--saved-sde-path" in values
        else None
    )
    data_config_path = (
        Path(values["--data-config-path"])
        if "--data-config-path" in values
        else None
    )
    downloads_manifest_path = (
        Path(values["--downloads-manifest-path"])
        if "--downloads-manifest-path" in values
        else None
    )
    dataset_id = values.get("--dataset-id")
    tokenizer_snapshot = (
        Path(values["--tokenizer-snapshot"])
        if "--tokenizer-snapshot" in values
        else None
    )
    benchmark_output = (
        Path(values["--benchmark-output"])
        if "--benchmark-output" in values
        else None
    )
    benchmark_metadata = (
        Path(values["--benchmark-metadata"])
        if "--benchmark-metadata" in values
        else None
    )
    benchmark_precision = values.get("--benchmark-precision")
    benchmark_fields = (benchmark_output, benchmark_metadata, benchmark_precision)
    if any(value is not None for value in benchmark_fields) and not all(
        value is not None for value in benchmark_fields
    ):
        raise ValueError("benchmark capture arguments must be provided together")
    if benchmark_output is not None:
        if not benchmark_output.is_absolute() or not benchmark_metadata.is_absolute():
            raise ValueError("benchmark paths must be absolute")
        if benchmark_precision != "author":
            raise ValueError("benchmark precision is invalid")
    if kind in {"langflow", "rdlm"} and expected_samples is None:
        raise ValueError(f"{kind} capture requires an expected sample count")
    if kind == "rdlm" and (saved_config_path is None or saved_sde_path is None):
        raise ValueError("RDLM capture requires saved config and SDE paths")
    generation_mode = values.get("--generation-mode", "unconditional")
    if generation_mode not in {"unconditional", "conditional_prefix"}:
        raise ValueError(f"unsupported generation mode: {generation_mode!r}")
    conditioning_manifest = (
        Path(values["--conditioning-manifest"])
        if "--conditioning-manifest" in values
        else None
    )

    def optional_int(name: str) -> int | None:
        text = values.get(name)
        if text is None:
            return None
        try:
            parsed = int(text)
        except ValueError as error:
            raise ValueError(f"{name} must be an integer") from error
        if parsed < 0:
            raise ValueError(f"{name} must be non-negative")
        return parsed

    conditional_fields = (
        conditioning_manifest,
        values.get("--conditioning-manifest-sha256"),
        values.get("--conditioning-config-sha256"),
        optional_int("--prefix-length"),
        optional_int("--evaluation-continuation-length"),
        optional_int("--prompt-count"),
        optional_int("--diversity-prompt-count"),
        optional_int("--completions-per-diversity-prompt"),
        values.get("--completion-schedule"),
    )
    if generation_mode == "conditional_prefix":
        if any(field is None for field in conditional_fields):
            raise ValueError("conditional capture requires the full conditioning contract")
        if conditioning_manifest is None or not conditioning_manifest.is_absolute():
            raise ValueError("conditional manifest path must be absolute")
    elif any(field is not None for field in conditional_fields):
        raise ValueError("unconditional capture must not include conditional fields")
    tokenizer_fields = (
        data_config_path,
        downloads_manifest_path,
        dataset_id,
        tokenizer_snapshot,
    )
    if kind in {"langflow", "rdlm"} and any(value is None for value in tokenizer_fields):
        raise ValueError(f"{kind} capture requires a locked tokenizer snapshot")
    return (
        CaptureInvocation(
            entrypoint=entrypoint,
            capture_path=capture_path,
            kind=kind,
            expected_samples=expected_samples,
            saved_config_path=saved_config_path,
            saved_sde_path=saved_sde_path,
            data_config_path=data_config_path,
            downloads_manifest_path=downloads_manifest_path,
            dataset_id=dataset_id,
            tokenizer_snapshot=tokenizer_snapshot,
            benchmark_output=benchmark_output,
            benchmark_metadata=benchmark_metadata,
            benchmark_precision=benchmark_precision,
            generation_mode=generation_mode,
            conditioning_manifest=conditioning_manifest,
            conditioning_manifest_sha256=values.get("--conditioning-manifest-sha256"),
            conditioning_config_sha256=values.get("--conditioning-config-sha256"),
            prefix_length=optional_int("--prefix-length"),
            evaluation_continuation_length=optional_int("--evaluation-continuation-length"),
            prompt_count=optional_int("--prompt-count"),
            diversity_prompt_count=optional_int("--diversity-prompt-count"),
            completions_per_diversity_prompt=optional_int("--completions-per-diversity-prompt"),
            completion_schedule=values.get("--completion-schedule"),
        ),
        forwarded,
    )


def _require_absolute_file(path: Path | None, label: str) -> Path:
    if path is None or not path.is_absolute() or path.is_symlink() or not path.is_file():
        raise ValueError(f"{label} is missing or unsafe: {path}")
    return path


def _load_tokenizer_binding(invocation: CaptureInvocation) -> TokenizerBinding:
    """Resolve one local snapshot against both immutable project data records."""

    config_path = _require_absolute_file(invocation.data_config_path, "data config")
    downloads_path = _require_absolute_file(
        invocation.downloads_manifest_path, "download manifest"
    )
    try:
        configuration = yaml.safe_load(config_path.read_text(encoding="utf-8"))
        downloads = json.loads(downloads_path.read_text(encoding="utf-8"))
        dataset = configuration["datasets"][invocation.dataset_id]
        tokenizer_id = dataset["tokenizer"]
        revision = configuration["models"][tokenizer_id]
        record = downloads["models"][tokenizer_id]
    except (OSError, UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as error:
        raise ValueError("locked tokenizer metadata is invalid") from error
    if downloads.get("schema_version") != 1:
        raise ValueError("download manifest schema is not supported")
    if not isinstance(tokenizer_id, str) or not isinstance(revision, str) or len(revision) != 40:
        raise ValueError("data config tokenizer binding is invalid")
    if record.get("repo_id") != tokenizer_id or record.get("revision") != revision:
        raise ValueError("download manifest tokenizer binding differs from data config")
    recorded_snapshot = record.get("snapshot_path")
    if not isinstance(recorded_snapshot, str) or not recorded_snapshot:
        raise ValueError("download manifest tokenizer snapshot is missing")
    snapshot_path = Path(recorded_snapshot)
    if not snapshot_path.is_absolute():
        snapshot_path = config_path.parents[1] / snapshot_path
    requested_snapshot = invocation.tokenizer_snapshot
    if requested_snapshot is None or not requested_snapshot.is_absolute():
        raise ValueError("tokenizer snapshot path must be absolute")
    if requested_snapshot != snapshot_path:
        raise ValueError("capture tokenizer snapshot differs from download manifest")
    parts = snapshot_path.parts
    try:
        recorded_revision = parts[parts.index("snapshots") + 1]
    except (ValueError, IndexError) as error:
        raise ValueError("download manifest tokenizer snapshot path is invalid") from error
    if recorded_revision != revision:
        raise ValueError("download manifest tokenizer snapshot revision differs from data config")
    if snapshot_path.is_symlink() or not snapshot_path.is_dir():
        raise ValueError(f"locked tokenizer snapshot is missing or unsafe: {snapshot_path}")
    return TokenizerBinding(tokenizer_id, revision, snapshot_path)


@contextmanager
def _offline_huggingface():
    names = ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
    previous = {name: os.environ.get(name) for name in names}
    os.environ.update({name: "1" for name in names})
    try:
        yield
    finally:
        for name, value in previous.items():
            if value is None:
                os.environ.pop(name, None)
            else:
                os.environ[name] = value


def _duo_canonical_logits_vocab_size(model: object) -> int | None:
    config = getattr(model, "config", None)
    model_type = getattr(config, "model_type", None)
    vocab_size = getattr(config, "vocab_size", None)
    if (
        isinstance(model_type, str)
        and model_type.lower() == "duo"
        and vocab_size == 50_258
    ):
        return 50_257
    return None


def _duo_canonical_sampler_vocab_size(owner: object) -> int | None:
    size = _duo_canonical_logits_vocab_size(getattr(owner, "backbone", None))
    current = getattr(owner, "vocab_size", None)
    if size is not None and current == size + 1:
        return size
    return None


def _adapt_hf_masked_lm_backbone(model: object) -> object:
    forward = getattr(model, "forward", None)
    if not callable(forward):
        return model
    try:
        parameters = inspect.signature(forward).parameters
    except (TypeError, ValueError):
        return model
    if {"x", "sigma"} <= set(parameters):
        return model
    if not {"input_ids", "timesteps"} <= set(parameters):
        return model
    try:
        import torch
    except ImportError:
        module_base = object
    else:
        module_base = torch.nn.Module

    class KeywordBackboneAdapter(module_base):
        def __init__(self, inner: object) -> None:
            if module_base is not object:
                super().__init__()
            self.inner = inner
            self._canonical_logits_vocab_size = _duo_canonical_logits_vocab_size(inner)

        def _canonical_logits(self, logits):
            size = self._canonical_logits_vocab_size
            shape = getattr(logits, "shape", ())
            if size is not None and shape and shape[-1] == size + 1:
                return logits[..., :size]
            return logits

        def forward(self, *args, **kwargs):
            if "x" in kwargs:
                kwargs["input_ids"] = kwargs.pop("x")
            elif args:
                kwargs["input_ids"] = args[0]
                args = args[1:]
            if "sigma" in kwargs:
                kwargs["timesteps"] = kwargs.pop("sigma")
            elif args:
                kwargs["timesteps"] = args[0]
                args = args[1:]
            kwargs.pop("class_cond", None)
            kwargs.pop("weights", None)
            call = self.inner if callable(self.inner) else self.inner.forward
            output = call(*args, **kwargs)
            if hasattr(output, "logits"):
                return self._canonical_logits(output.logits)
            if isinstance(output, tuple) and output:
                return self._canonical_logits(output[0])
            return self._canonical_logits(output)

        def __getattr__(self, name: str) -> object:
            try:
                return super().__getattr__(name)
            except AttributeError:
                return getattr(self.inner, name)

    if module_base is object:
        KeywordBackboneAdapter.__call__ = KeywordBackboneAdapter.forward
    return KeywordBackboneAdapter(model)


@contextmanager
def _patched_hf_masked_lm_backbone():
    try:
        transformers = importlib.import_module("transformers")
    except ImportError:
        yield
        return
    auto_model = getattr(transformers, "AutoModelForMaskedLM", None)
    original = getattr(auto_model, "from_pretrained", None)
    if not callable(original):
        yield
        return

    def from_pretrained(*args, **kwargs):
        return _adapt_hf_masked_lm_backbone(original(*args, **kwargs))

    auto_model.from_pretrained = from_pretrained
    try:
        yield
    finally:
        auto_model.from_pretrained = original


def _load_entrypoint(path: Path) -> ModuleType:
    sys.path.insert(0, str(path.parent))
    module_name = "dlb_pinned_upstream_main"
    specification = importlib.util.spec_from_file_location(module_name, path)
    if specification is None or specification.loader is None:
        raise ValueError(f"cannot load upstream entrypoint: {path}")
    module = importlib.util.module_from_spec(specification)
    previous = sys.modules.get(module_name)
    sys.modules[module_name] = module
    try:
        specification.loader.exec_module(module)
    except BaseException:
        if previous is None:
            sys.modules.pop(module_name, None)
        else:
            sys.modules[module_name] = previous
        raise
    return module


def _sample_owner(module: ModuleType) -> type:
    diffusion = getattr(module, "diffusion", None)
    if diffusion is not None and hasattr(diffusion, "Diffusion"):
        return diffusion.Diffusion
    algo = getattr(module, "algo", None)
    trainer_base = getattr(algo, "trainer_base", None)
    owner = getattr(trainer_base, "TrainerBase", None)
    if isinstance(owner, type):
        return owner
    raise ValueError("upstream entrypoint has no supported sampling owner")


def _write_capture(
    path: Path,
    texts: list[object],
    token_rows: list[object],
    expected: int | None,
    extras: list[dict[str, object]] | None = None,
) -> None:
    if len(texts) != len(token_rows):
        raise ValueError("upstream token and text sample counts differ")
    if extras is not None and len(extras) != len(token_rows):
        raise ValueError("upstream conditional metadata and token sample counts differ")
    if expected is not None and len(texts) != expected:
        raise ValueError(f"expected {expected} captured samples, found {len(texts)}")
    samples: list[dict[str, object]] = []
    for index, (text, tokens) in enumerate(zip(texts, token_rows)):
        if not isinstance(text, str):
            raise ValueError(f"captured text {index} is not a string")
        if not isinstance(tokens, list):
            raise ValueError(f"captured token row {index} is not a list")
        payload = {"sample_id": index, "text": text, "token_ids": tokens}
        if extras is not None:
            payload.update(extras[index])
        samples.append(payload)
    atomic_json_write(path, {"schema": "dlb-upstream-token-capture-v1", "samples": samples})


def _rows(value: object, *, label: str = "token rows") -> list[object]:
    if hasattr(value, "detach"):
        value = value.detach().cpu().tolist()
    if not isinstance(value, list):
        raise ValueError(f"{label} returned invalid token rows")
    rows: list[object] = []
    for row in value:
        if hasattr(row, "detach"):
            row = row.detach().cpu().tolist()
        elif not isinstance(row, list):
            row = list(row)
        rows.append(row)
    return rows


def _run_main(
    module: ModuleType,
    entrypoint: Path,
    forwarded: list[str],
    *,
    hydra_config_path: bool = True,
) -> None:
    previous_argv = sys.argv
    upstream_arguments = (
        _forwarded_with_hydra_config_path(entrypoint, forwarded)
        if hydra_config_path
        else forwarded
    )
    try:
        sys.argv = [str(entrypoint), *upstream_arguments]
        module.main()
    finally:
        sys.argv = previous_argv


def _forwarded_with_hydra_config_path(entrypoint: Path, forwarded: list[str]) -> list[str]:
    if any(
        argument in {"--config-path", "-cp"}
        or argument.startswith("--config-path=")
        or argument.startswith("-cp=")
        for argument in forwarded
    ):
        return forwarded
    config_dir = entrypoint.parent / "configs"
    if config_dir.is_dir() and not config_dir.is_symlink():
        return [f"--config-path={config_dir.resolve()}", *forwarded]
    return forwarded


def _capture_teacher(
    module: ModuleType, invocation: CaptureInvocation, forwarded: list[str]
) -> None:
    owner = _sample_owner(module)
    original = owner.restore_model_and_sample
    captured: list[dict[str, object]] = []

    def current_conditioning_batch(self):
        if invocation.generation_mode != "conditional_prefix":
            return None
        if (
            invocation.conditioning_manifest is None
            or invocation.conditioning_manifest_sha256 is None
            or invocation.prompt_count is None
            or invocation.diversity_prompt_count is None
            or invocation.completions_per_diversity_prompt is None
        ):
            raise ValueError("conditional capture invocation is incomplete")
        offset = len(captured)
        if offset < invocation.prompt_count:
            completion_id = 0
            prompt_start = offset
            limit = invocation.prompt_count
        else:
            diversity_offset = offset - invocation.prompt_count
            completion_id = diversity_offset // invocation.diversity_prompt_count + 1
            prompt_start = diversity_offset % invocation.diversity_prompt_count
            limit = invocation.diversity_prompt_count
        if completion_id >= invocation.completions_per_diversity_prompt:
            raise ValueError("conditional capture schedule is exhausted")
        batch_size = int(getattr(self.config.loader, "eval_batch_size"))
        current = min(batch_size, limit - prompt_start)
        if current != batch_size:
            raise ValueError("conditional batch size crosses a schedule boundary")
        vocab_size = int(getattr(self, "vocab_size", 0) or len(self.tokenizer))
        return load_conditioning_batch(
            invocation.conditioning_manifest,
            invocation.conditioning_manifest_sha256,
            completion_id=completion_id,
            prompt_start=prompt_start,
            batch_size=batch_size,
            device=str(self.device),
            vocab_size=vocab_size,
        )

    def install_conditioning(self, batch):
        if batch is None:
            return ExitStack()
        family = invocation.entrypoint.parent.name
        stack = ExitStack()
        if family == "mdlm":
            project = token_project_fn(batch.prefix_token_ids)
            if hasattr(self, "_sample_prior"):
                original_prior = self._sample_prior

                def prior(*args, **kwargs):
                    return project(original_prior(*args, **kwargs).to(self.device))

                stack.enter_context(patched_attribute(self, "_sample_prior", prior))
            for name in ("_ddpm_update", "_ddpm_caching_update", "_analytic_update", "_denoiser_update"):
                original_update = getattr(self, name, None)
                if not callable(original_update):
                    continue

                def update(*args, __original=original_update, **kwargs):
                    if "x" in kwargs:
                        kwargs["x"] = project(kwargs["x"])
                    elif args:
                        args = (project(args[0]), *args[1:])
                    output = __original(*args, **kwargs)
                    if isinstance(output, tuple):
                        return (*output[:-1], project(output[-1]))
                    return project(output)

                stack.enter_context(patched_attribute(self, name, update))
        elif family == "duo":
            sampler_vocab_size = _duo_canonical_sampler_vocab_size(self)
            if sampler_vocab_size is not None:
                stack.enter_context(patched_attribute(self, "vocab_size", sampler_vocab_size))
            project = token_project_fn(batch.prefix_token_ids)
            if hasattr(self, "prior_sample"):
                original_prior = self.prior_sample

                def prior(*args, **kwargs):
                    return project(original_prior(*args, **kwargs).to(self.device))

                stack.enter_context(patched_attribute(self, "prior_sample", prior))
            for name in ("_ancestral_update", "_analytic_update", "_denoiser_update"):
                original_update = getattr(self, name, None)
                if not callable(original_update):
                    continue

                def update(*args, __original=original_update, **kwargs):
                    if "x" in kwargs:
                        kwargs["x"] = project(kwargs["x"])
                    elif args:
                        args = (project(args[0]), *args[1:])
                    output = __original(*args, **kwargs)
                    if isinstance(output, tuple):
                        return (*output[:-1], project(output[-1]))
                    return project(output)

                stack.enter_context(patched_attribute(self, name, update))
        elif family == "flm":
            project = vocab_project_fn(batch.prefix_token_ids)
            if hasattr(self, "prior_sample"):
                original_prior = self.prior_sample

                def prior(*args, **kwargs):
                    return project(original_prior(*args, **kwargs).to(self.device))

                stack.enter_context(patched_attribute(self, "prior_sample", prior))
            original_forward = self.forward

            def forward(*args, **kwargs):
                if "xt" in kwargs:
                    kwargs["xt"] = project(kwargs["xt"])
                elif args:
                    args = (project(args[0]), *args[1:])
                return original_forward(*args, **kwargs)

            stack.enter_context(patched_attribute(self, "forward", forward))
        elif family == "candi":
            original_generate = self.generate_samples
            original_generate_sample_prompt = self.generate_sample_prompt
            stack.enter_context(
                patched_attribute(
                    self,
                    "generate_sample_prompt",
                    adapt_candi_generate_sample_prompt(original_generate_sample_prompt),
                )
            )

            def generate(*args, **kwargs):
                prompt_tokens, prompt_mask = candi_prompt_conditioning(
                    batch.prefix_token_ids,
                    sequence_length=int(getattr(self, "num_tokens")),
                )
                kwargs["prompt_tokens"] = prompt_tokens
                kwargs["prompt_mask"] = prompt_mask
                return original_generate(*args, **kwargs)

            stack.enter_context(patched_attribute(self, "generate_samples", generate))
        else:
            raise ValueError(f"unsupported conditional teacher family: {family}")
        return stack

    def capture(self, *args, **kwargs):
        if invocation.benchmark_output is not None:
            benchmark_stack = ExitStack()
            if (
                invocation.generation_mode == "conditional_prefix"
                and invocation.entrypoint.parent.name == "duo"
            ):
                sampler_vocab_size = _duo_canonical_sampler_vocab_size(self)
                if sampler_vocab_size is not None:
                    benchmark_stack.enter_context(
                        patched_attribute(self, "vocab_size", sampler_vocab_size)
                    )

            with benchmark_stack:
                num_steps = kwargs.get("num_steps", args[0] if args else None)
                eps = kwargs.get("eps", 1e-5)
                if type(num_steps) is not int or num_steps <= 0:
                    raise ValueError("benchmark sampler requires a positive num_steps")
                if invocation.entrypoint.parent.name == "mdlm":
                    if self.ema:
                        parameters = tuple(
                            list(self.backbone.parameters()) + list(self.noise.parameters())
                        )
                        self.ema.store(parameters)
                        self.ema.copy_to(parameters)
                    self.backbone.eval()
                    self.noise.eval()
                    generate_one = lambda: self._sample(num_steps=num_steps, eps=eps)
                else:
                    self._eval_mode()
                    generate_one = lambda: self.generate_samples(
                        num_samples=1, num_steps=num_steps, eps=eps
                    )
                try:
                    return benchmark_and_publish(
                        generate_one,
                        model=self,
                        output=invocation.benchmark_output,
                        metadata_path=invocation.benchmark_metadata,
                        precision=invocation.benchmark_precision,
                    )
                finally:
                    if invocation.entrypoint.parent.name == "mdlm":
                        if self.ema:
                            self.ema.restore(parameters)
                        self.backbone.train()
                        self.noise.train()
                    else:
                        self._train_mode()
        batch = current_conditioning_batch(self)
        with install_conditioning(self, batch):
            result = original(self, *args, **kwargs)
        if batch is not None:
            if isinstance(result, list) and result and all(hasattr(row, "shape") for row in result):
                import torch

                result = torch.stack(result)
            result = clamp_token_prefix(result, batch.prefix_token_ids)
        token_rows = _rows(result, label="teacher sampler")
        text_rows = list(self.tokenizer.batch_decode(result))
        if len(token_rows) != len(text_rows):
            raise ValueError("upstream token and text batch sizes differ")
        for batch_index, (tokens, text) in enumerate(zip(token_rows, text_rows)):
            payload = {"sample_id": len(captured), "text": text, "token_ids": tokens}
            if batch is not None:
                prefix = batch.prefix_token_ids[batch_index].detach().cpu().tolist()
                reference = batch.reference_token_ids[batch_index].detach().cpu().tolist()
                if tokens[: len(prefix)] != prefix:
                    raise ValueError("conditional sampler returned a prefix mismatch")
                payload.update(
                    {
                        "prompt_id": batch.prompt_ids[batch_index],
                        "completion_id": batch.completion_id,
                        "source_index": batch.source_indices[batch_index],
                        "prefix_token_ids": prefix,
                        "reference_token_ids": reference,
                        "full_token_ids": tokens,
                        "prefix_text": self.tokenizer.decode(prefix),
                        "continuation_text": self.tokenizer.decode(tokens[len(prefix) :]),
                        "reference_text": self.tokenizer.decode(reference),
                        "full_text": text,
                        "prefix_exact_match": True,
                    }
                )
            captured.append(payload)
        return result

    owner.restore_model_and_sample = capture
    with _patched_hf_masked_lm_backbone():
        _run_main(module, invocation.entrypoint, forwarded)
    atomic_json_write(
        invocation.capture_path,
        {"schema": "dlb-upstream-token-capture-v1", "samples": captured},
    )


def _capture_langflow(
    module: ModuleType,
    invocation: CaptureInvocation,
    forwarded: list[str],
    tokenizer: TokenizerBinding,
) -> None:
    owner = getattr(module, "LangFlow", None)
    auto_tokenizer = getattr(module, "AutoTokenizer", None)
    if not isinstance(owner, type) or auto_tokenizer is None:
        raise ValueError("LangFlow entrypoint lacks its model or tokenizer")
    original_generate = owner.generate_samples
    token_rows: list[object] = []
    texts: list[object] = []
    conditional_extras: list[dict[str, object]] = []

    def generate(self, *args, **kwargs):
        conditioning = None
        project_embeddings = None
        original_forward = None
        if invocation.generation_mode == "conditional_prefix":
            if invocation.conditioning_manifest is None or invocation.conditioning_manifest_sha256 is None:
                raise ValueError("LangFlow conditional capture requires a prompt manifest")
            batch_size = int(kwargs.get("num_samples", 1))
            conditioning = load_conditioning_batch(
                invocation.conditioning_manifest,
                invocation.conditioning_manifest_sha256,
                completion_id=0 if len(token_rows) < (invocation.prompt_count or 0) else (len(token_rows) - (invocation.prompt_count or 0)) // (invocation.diversity_prompt_count or 1) + 1,
                prompt_start=len(token_rows) if len(token_rows) < (invocation.prompt_count or 0) else (len(token_rows) - (invocation.prompt_count or 0)) % (invocation.diversity_prompt_count or 1),
                batch_size=batch_size,
                device=str(next(self.parameters()).device),
                vocab_size=int(self.config.vocab_size),
            )
            import torch

            clean_one_hot = torch.nn.functional.one_hot(
                conditioning.prefix_token_ids, num_classes=int(self.config.vocab_size)
            ).to(dtype=next(self.parameters()).dtype, device=conditioning.prefix_token_ids.device)
            clean_embeddings = self._embed_tokens(clean_one_hot)
            project_embeddings = embedding_project_fn(clean_embeddings)
            original_forward = self.forward

            def forward(*args, **kwargs):
                if "noisy_embeds" in kwargs and kwargs["noisy_embeds"] is not None:
                    kwargs["noisy_embeds"] = project_embeddings(kwargs["noisy_embeds"])
                if "x_self_cond" in kwargs and kwargs["x_self_cond"] is not None:
                    kwargs["x_self_cond"] = project_embeddings(kwargs["x_self_cond"])
                return original_forward(*args, **kwargs)

        with (
            patched_attribute(self, "forward", forward)
            if project_embeddings is not None and original_forward is not None
            else nullcontext()
        ):
            if invocation.benchmark_output is None:
                result = original_generate(self, *args, **kwargs)
            else:
                benchmark_kwargs = dict(kwargs)
                benchmark_kwargs["num_samples"] = 1
                result = benchmark_and_publish(
                    lambda: original_generate(self, *args, **benchmark_kwargs),
                    model=self,
                    output=invocation.benchmark_output,
                    metadata_path=invocation.benchmark_metadata,
                    precision=invocation.benchmark_precision,
                )
        if conditioning is not None:
            result = clamp_token_prefix(result, conditioning.prefix_token_ids)
        rows = result.detach().cpu().tolist()
        if not isinstance(rows, list):
            raise ValueError("LangFlow sampler returned invalid token rows")
        if conditioning is not None:
            for batch_index, row in enumerate(rows):
                prefix = conditioning.prefix_token_ids[batch_index].detach().cpu().tolist()
                reference = conditioning.reference_token_ids[batch_index].detach().cpu().tolist()
                if row[: len(prefix)] != prefix:
                    raise ValueError("LangFlow conditional sampler returned a prefix mismatch")
                conditional_extras.append(
                    {
                        "prompt_id": conditioning.prompt_ids[batch_index],
                        "completion_id": conditioning.completion_id,
                        "source_index": conditioning.source_indices[batch_index],
                        "prefix_token_ids": prefix,
                        "reference_token_ids": reference,
                        "full_token_ids": row,
                        "prefix_exact_match": True,
                    }
                )
        token_rows.extend(rows)
        return result

    class TokenizerProxy:
        def __init__(self, tokenizer: object) -> None:
            self._tokenizer = tokenizer

        def __getattr__(self, name: str) -> object:
            return getattr(self._tokenizer, name)

        def batch_decode(self, *args, **kwargs):
            decoded = list(self._tokenizer.batch_decode(*args, **kwargs))
            texts[:] = decoded
            return decoded

    class AutoTokenizerProxy:
        @staticmethod
        def from_pretrained(*args, **kwargs):
            del args
            kwargs.pop("revision", None)
            kwargs["local_files_only"] = True
            loaded = auto_tokenizer.from_pretrained(str(tokenizer.snapshot), **kwargs)
            return TokenizerProxy(loaded)

    owner.generate_samples = generate
    module.AutoTokenizer = AutoTokenizerProxy
    with _offline_huggingface():
        parameters = inspect.signature(_run_main).parameters
        if "hydra_config_path" in parameters:
            _run_main(module, invocation.entrypoint, forwarded, hydra_config_path=False)
        else:
            _run_main(module, invocation.entrypoint, forwarded)
    _write_capture(
        invocation.capture_path,
        texts,
        token_rows,
        invocation.expected_samples,
        conditional_extras if invocation.generation_mode == "conditional_prefix" else None,
    )


def _model_path(forwarded: list[str]) -> Path:
    prefix = "model_path="
    values = [argument[len(prefix) :] for argument in forwarded if argument.startswith(prefix)]
    if len(values) != 1 or not values[0]:
        raise ValueError("RDLM capture requires one model_path override")
    path = Path(values[0])
    if not path.is_absolute() or path.is_symlink() or not path.is_file():
        raise ValueError(f"RDLM checkpoint is missing or unsafe: {path}")
    return path


def _positive_override(forwarded: list[str], key: str) -> int:
    prefix = key + "="
    values = [argument[len(prefix) :] for argument in forwarded if argument.startswith(prefix)]
    if len(values) != 1:
        raise ValueError(f"RDLM capture requires one {key} override")
    try:
        value = int(values[0])
    except ValueError as error:
        raise ValueError(f"RDLM {key} must be an integer") from error
    if value <= 0:
        raise ValueError(f"RDLM {key} must be positive")
    return value


def _require_saved_file(path: Path | None, label: str) -> Path:
    if (
        path is None
        or not path.is_absolute()
        or path.is_symlink()
        or not path.is_file()
        or path.stat().st_size <= 0
    ):
        raise ValueError(f"saved RDLM {label} is missing or unsafe: {path}")
    return path


def _rdlm_benchmark_sampling_fn(
    sampling_fn: Callable[[object], object], invocation: CaptureInvocation
) -> Callable[[object], object]:
    """Publish one benchmark the first time RDLM reaches its loaded sampling_fn."""

    if invocation.benchmark_output is None:
        return sampling_fn
    published = False

    def sample(model: object) -> object:
        nonlocal published
        if published:
            return sampling_fn(model)
        published = True
        return benchmark_and_publish(
            lambda: sampling_fn(model),
            model=model,
            output=invocation.benchmark_output,
            metadata_path=invocation.benchmark_metadata,
            precision=invocation.benchmark_precision,
        )

    return sample


def _capture_rdlm(
    module: ModuleType, invocation: CaptureInvocation, forwarded: list[str]
) -> None:
    from omegaconf import OmegaConf

    tokenizer = _load_tokenizer_binding(invocation)
    saved_config_path = _require_saved_file(invocation.saved_config_path, "config")
    saved_sde_path = _require_saved_file(invocation.saved_sde_path, "SDE")
    checkpoint_path = _model_path(forwarded)
    expected = invocation.expected_samples
    if type(expected) is not int or expected <= 0:
        raise ValueError("RDLM capture requires a positive expected sample count")
    batch_size = _positive_override(forwarded, "sampling.batch_per_gpu")
    sample_batches = math.ceil(expected / batch_size)
    saved_config = OmegaConf.load(saved_config_path)
    if OmegaConf.select(saved_config, "model.length") != 128:
        raise ValueError("saved RDLM config does not declare LM1B model length 128")
    if OmegaConf.select(saved_config, "data.train") != "lm1b":
        raise ValueError("saved RDLM config is not the LM1B training config")

    run_sample = importlib.import_module("run_sample")
    data_module = getattr(run_sample, "data", None)
    transformers_module = getattr(data_module, "transformers", None)
    original_bert_tokenizer = getattr(transformers_module, "BertTokenizer", None)
    original_auto_tokenizer = getattr(data_module, "AutoTokenizer", None)
    if original_bert_tokenizer is None or original_auto_tokenizer is None:
        raise ValueError("RDLM data module lacks its tokenizer constructors")
    original_torch = run_sample.torch
    original_open: Callable[..., object] = builtins.open
    original_tqdm = run_sample.tqdm
    original_instantiate = run_sample.instantiate
    original_shift_factory = run_sample.sutils.find_bos_and_shift_fn
    original_sampling_factory = (
        run_sample.sampling.get_sampling_fn
        if invocation.benchmark_output is not None
        or invocation.generation_mode == "conditional_prefix"
        else None
    )
    texts: list[object] = []
    token_rows: list[object] = []
    conditional_extras: list[dict[str, object]] = []
    saved_config_used = False
    saved_sde_opened = False
    saved_sde_used = False

    def locked_tokenizer_proxy(original):
        class LocalTokenizerProxy:
            @staticmethod
            def from_pretrained(*args, **kwargs):
                del args
                kwargs.pop("revision", None)
                kwargs["local_files_only"] = True
                return original.from_pretrained(str(tokenizer.snapshot), **kwargs)

        return LocalTokenizerProxy

    class TorchProxy:
        def __getattr__(self, name: str) -> object:
            return getattr(original_torch, name)

        def load(self, path, *args, **kwargs):
            nonlocal saved_config_used
            state = original_torch.load(path, *args, **kwargs)
            if Path(path).resolve() != checkpoint_path.resolve():
                return state
            if not isinstance(state, dict):
                raise ValueError("RDLM checkpoint state is not a mapping")
            saved_config_used = True
            return {**state, "config": saved_config}

    def routed_open(path, mode="r", *args, **kwargs):
        nonlocal saved_sde_opened
        requested = Path(path)
        if "r" in mode and "b" in mode and requested.name == "sde.pkl":
            saved_sde_opened = True
            return original_open(saved_sde_path, mode, *args, **kwargs)
        return original_open(path, mode, *args, **kwargs)

    def instantiate(config, *args, **kwargs):
        nonlocal saved_sde_used
        is_sde = {"manifold", "scheduler", "prior_dist"}.issubset(kwargs)
        if is_sde and "preprocessed" not in kwargs:
            raise ValueError("saved SDE preprocessing was not used by RDLM")
        if is_sde:
            saved_sde_used = True
        return original_instantiate(config, *args, **kwargs)

    def exact_sample_batches(iterable, *args, **kwargs):
        del iterable
        return original_tqdm(range(sample_batches), *args, **kwargs)

    def rdlm_conditioning_batch(batch_index: int, device: str):
        if invocation.generation_mode != "conditional_prefix":
            return None
        if (
            invocation.conditioning_manifest is None
            or invocation.conditioning_manifest_sha256 is None
            or invocation.prompt_count is None
            or invocation.diversity_prompt_count is None
            or invocation.completions_per_diversity_prompt is None
        ):
            raise ValueError("RDLM conditional capture invocation is incomplete")
        offset = batch_index * batch_size
        if offset < invocation.prompt_count:
            completion_id = 0
            prompt_start = offset
            limit = invocation.prompt_count
        else:
            diversity_offset = offset - invocation.prompt_count
            completion_id = diversity_offset // invocation.diversity_prompt_count + 1
            prompt_start = diversity_offset % invocation.diversity_prompt_count
            limit = invocation.diversity_prompt_count
        if completion_id >= invocation.completions_per_diversity_prompt:
            raise ValueError("RDLM conditional capture schedule is exhausted")
        if min(batch_size, limit - prompt_start) != batch_size:
            raise ValueError("RDLM conditional batch size crosses a schedule boundary")
        original_tokens = OmegaConf.select(saved_config, "original_tokens")
        if type(original_tokens) is not int or original_tokens <= 0:
            raise ValueError("saved RDLM config does not declare original_tokens")
        return load_conditioning_batch(
            invocation.conditioning_manifest,
            invocation.conditioning_manifest_sha256,
            completion_id=completion_id,
            prompt_start=prompt_start,
            batch_size=batch_size,
            device=device,
            vocab_size=original_tokens,
        )

    def sampling_factory(*args, **kwargs):
        if original_sampling_factory is None:
            raise ValueError("RDLM benchmark sampling factory is unavailable")
        device = kwargs.get("device", args[4] if len(args) > 4 else "cpu")
        active_project = None
        batch_index = 0
        if invocation.generation_mode == "conditional_prefix":
            kwargs = dict(kwargs)

            def project(state):
                return state if active_project is None else active_project(state)

            kwargs["proj_fn"] = project
        sampling_fn = original_sampling_factory(*args, **kwargs)
        sampling_fn = _rdlm_benchmark_sampling_fn(sampling_fn, invocation)

        def sample(model: object) -> object:
            nonlocal active_project, batch_index
            batch = rdlm_conditioning_batch(batch_index, str(device))
            if batch is not None:
                token_size = OmegaConf.select(saved_config, "tokens")
                original_tokens = OmegaConf.select(saved_config, "original_tokens")
                if (
                    type(token_size) is not int
                    or token_size <= 1
                    or type(original_tokens) is not int
                    or original_tokens <= 0
                ):
                    raise ValueError("saved RDLM config has invalid base token contract")
                digits_per_token = math.ceil(math.log(original_tokens) / math.log(token_size))
                active_project = rdlm_project_fn(
                    batch.prefix_token_ids,
                    base=token_size,
                    digits_per_token=digits_per_token,
                )
                for row_index in range(batch.prefix_token_ids.shape[0]):
                    conditional_extras.append(
                        {
                            "prompt_id": batch.prompt_ids[row_index],
                            "completion_id": batch.completion_id,
                            "source_index": batch.source_indices[row_index],
                            "prefix_token_ids": batch.prefix_token_ids[row_index]
                            .detach()
                            .cpu()
                            .tolist(),
                            "reference_token_ids": batch.reference_token_ids[row_index]
                            .detach()
                            .cpu()
                            .tolist(),
                        }
                    )
                batch_index += 1
            try:
                return sampling_fn(model)
            finally:
                active_project = None

        return sample

    def capture_shift_factory(*args, **kwargs):
        shift = original_shift_factory(*args, **kwargs)
        tokenizer_object = args[2] if len(args) >= 3 else kwargs.get("tokenizer")

        def capture_shift(samples):
            sentences, shifted = shift(samples)
            texts[:] = list(sentences)
            token_rows[:] = _rows(shifted)
            if invocation.generation_mode == "conditional_prefix":
                if len(conditional_extras) < len(token_rows):
                    raise ValueError("RDLM conditional metadata is shorter than generated samples")
                for index, row in enumerate(token_rows):
                    prefix = conditional_extras[index]["prefix_token_ids"]
                    reference = conditional_extras[index]["reference_token_ids"]
                    if not isinstance(prefix, list) or row[: len(prefix)] != prefix:
                        raise ValueError("RDLM conditional sampler returned a prefix mismatch")
                    conditional_extras[index].update(
                        {
                            "full_token_ids": row,
                            "prefix_exact_match": True,
                            "prefix_text": (
                                tokenizer_object.decode(prefix)
                                if tokenizer_object is not None
                                else " ".join(map(str, prefix))
                            ),
                            "continuation_text": (
                                tokenizer_object.decode(row[len(prefix) :])
                                if tokenizer_object is not None
                                else " ".join(map(str, row[len(prefix) :]))
                            ),
                            "reference_text": (
                                tokenizer_object.decode(reference)
                                if tokenizer_object is not None
                                else " ".join(map(str, reference))
                            ),
                            "full_text": texts[index],
                        }
                    )
            return sentences, shifted

        return capture_shift

    def inline_spawn(function, args=(), nprocs=1, join=True, **kwargs):
        del join, kwargs
        if nprocs != 1:
            raise ValueError("RDLM capture requires ngpus=1")
        return function(0, *args)

    run_sample.torch = TorchProxy()
    data_module.transformers.BertTokenizer = locked_tokenizer_proxy(
        original_bert_tokenizer
    )
    data_module.AutoTokenizer = locked_tokenizer_proxy(original_auto_tokenizer)
    run_sample.open = routed_open
    run_sample.tqdm = exact_sample_batches
    if original_sampling_factory is not None:
        run_sample.sampling.get_sampling_fn = sampling_factory
    run_sample.instantiate = instantiate
    run_sample.sutils.find_bos_and_shift_fn = capture_shift_factory
    module.mp.spawn = inline_spawn
    with _offline_huggingface():
        _run_main(module, invocation.entrypoint, forwarded)
    if not saved_config_used:
        raise ValueError("saved RDLM config was not used")
    if not saved_sde_opened or not saved_sde_used:
        raise ValueError("saved RDLM SDE was not used")
    if len(texts) < expected or len(token_rows) < expected:
        raise ValueError(
            f"expected {expected} captured samples, found {min(len(texts), len(token_rows))}"
        )
    _write_capture(
        invocation.capture_path,
        texts[:expected],
        token_rows[:expected],
        expected,
        conditional_extras[:expected]
        if invocation.generation_mode == "conditional_prefix"
        else None,
    )


def main(argv: list[str] | None = None) -> int:
    invocation, forwarded = _split_arguments(
        list(sys.argv[1:] if argv is None else argv)
    )
    module = _load_entrypoint(invocation.entrypoint)
    if invocation.kind == "langflow":
        tokenizer = _load_tokenizer_binding(invocation)
        _capture_langflow(module, invocation, forwarded, tokenizer)
    elif invocation.kind == "rdlm":
        _capture_rdlm(module, invocation, forwarded)
    else:
        _capture_teacher(module, invocation, forwarded)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())