"""Adapter for the pinned CANDI hybrid discrete/continuous sampler."""

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


class CANDIAdapter(BaseTeacherAdapter):
    identity = "dlb.adapters.candi:v1"
    upstream = "candi"
    supported_models = frozenset({"candi"})
    teacher_families = {"candi": "hybrid_candi"}
    batch_sizes = {("candi", "lm1b"): 2, ("candi", "owt"): 2}

    def _loader_batch_size(self, request: RunRequest) -> int:
        del request
        return 16

    def _sampling_overrides(
        self, request: RunRequest, checkpoint: CheckpointSelection
    ) -> list[str]:
        del request, checkpoint
        return [
            "algo=candi",
            "algo.sampler=cached",
            "algo.mixed_coeff=0.5",
            "algo.step_size=1.0",
            "algo.temp=1.0",
            "algo.use_percentile_scheduling=True",
            "sampling.noise_removal=ancestral",
            "trainer.accumulate_grad_batches=1",
            "eval.disable_ema=False",
        ]


adapter = CANDIAdapter()