File size: 1,345 Bytes
6b56858 | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 | """Run Duo+DCD with a narrow compatibility patch for the pinned Duo source."""
from __future__ import annotations
from pathlib import Path
import sys
from typing import Any
def _patch_distillation_nll() -> None:
import algo
original = algo.Distillation.nll
if getattr(original, "_dlb_accepts_labels", False):
return
def nll(
self: Any,
x0: Any,
labels_or_output_tokens: Any = None,
output_tokens: Any = None,
current_accumulation_step: Any = None,
train_mode: Any = None,
) -> Any:
if output_tokens is None:
output_tokens = labels_or_output_tokens
return original(
self,
x0,
output_tokens,
current_accumulation_step,
train_mode,
)
nll._dlb_accepts_labels = True # type: ignore[attr-defined]
algo.Distillation.nll = nll
def main() -> int:
import hydra
sys.path.insert(0, str(Path.cwd()))
_patch_distillation_nll()
import main as upstream_main
@hydra.main(
version_base=None,
config_path="../upstreams/duo/configs",
config_name="config",
)
def run(config: Any) -> None:
upstream_main.main.__wrapped__(config)
return int(run() or 0)
if __name__ == "__main__":
raise SystemExit(main())
|