| """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 |
| 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()) |
|
|
