| """Adapter for the pinned masked discrete language model.""" | |
| from dlb.adapters.base import BaseTeacherAdapter, CheckpointSelection | |
| from dlb.runner import RunRequest | |
| class MDLMAdapter(BaseTeacherAdapter): | |
| identity = "dlb.adapters.mdlm:v1" | |
| upstream = "mdlm" | |
| supported_models = frozenset({"mdlm"}) | |
| teacher_families = {"mdlm": "masked_mdlm"} | |
| batch_sizes = {("mdlm", "lm1b"): 16, ("mdlm", "owt"): 16} | |
| def _data_config(self, request: RunRequest) -> str: | |
| return "lm1b" if request.dataset_id == "lm1b" else "openwebtext-split" | |
| def _sampling_overrides( | |
| self, request: RunRequest, checkpoint: CheckpointSelection | |
| ) -> list[str]: | |
| del checkpoint | |
| backbone = "hf_dit" if request.dataset_id == "owt" else "dit" | |
| return [ | |
| f"backbone={backbone}", | |
| "parameterization=subs", | |
| "sampling.predictor=ddpm", | |
| "sampling.noise_removal=True", | |
| "eval.disable_ema=False", | |
| ] | |
| adapter = MDLMAdapter() | |