"""Adapter for pinned Duo and discrete-consistency-distilled Duo."""

from dlb.adapters.base import BaseTeacherAdapter, CheckpointSelection
from dlb.runner import RunRequest


class DuoAdapter(BaseTeacherAdapter):
    identity = "dlb.adapters.duo:v1"
    upstream = "duo"
    supported_models = frozenset({"duo", "duo_dcd"})
    teacher_families = {"duo": "uniform_duo", "duo_dcd": "uniform_duo"}
    batch_sizes = {
        ("duo", "lm1b"): 64,
        ("duo", "owt"): 8,
        ("duo_dcd", "lm1b"): 64,
        ("duo_dcd", "owt"): 8,
    }

    def _sampling_overrides(
        self, request: RunRequest, checkpoint: CheckpointSelection
    ) -> list[str]:
        del checkpoint
        backbone = "hf_dit" if request.dataset_id == "owt" else "dit"
        return [
            "algo=duo_base",
            f"algo.backbone={backbone}",
            "sampling.predictor=ancestral",
            "sampling.noise_removal=ancestral",
            "eval.disable_ema=False",
        ]


adapter = DuoAdapter()