File size: 1,011 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 | """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()
|