File size: 1,005 Bytes
784f572 | 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 | """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()
|