-
Notifications
You must be signed in to change notification settings - Fork 115
Expand file tree
/
Copy pathtrain_outer.py
More file actions
59 lines (48 loc) · 2.15 KB
/
Copy pathtrain_outer.py
File metadata and controls
59 lines (48 loc) · 2.15 KB
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
#!/usr/bin/env python3
"""Unified outer-loop trainer for RecursiveMAS.
Pick the collaboration style with --style; everything after it is forwarded to that style's
outer trainer unchanged. The recursive gradient path through the (frozen) inner adapters is
ALWAYS preserved, so the outer RecursiveLink modules are trained correctly across rounds.
python train_outer.py --style sequential_scaled --agent1_model_name_or_path ... ...
python train_outer.py --style mixture --agent1_model_name_or_path ... ...
python train_outer.py --style distillation --expert_model_name_or_path ... ...
python train_outer.py --style deliberation --reflector_model_name_or_path ... ...
Run `python train_outer.py --style <style> --help` to see that style's full argument list.
"""
import argparse
import sys
# Style -> outer family. sequential_light and sequential_scaled both use the sequential family
# (same recursive loop + roles); they differ only in which base models you pass.
STYLE_TO_FAMILY = {
"sequential_light": "sequential",
"sequential_scaled": "sequential",
"mixture": "hie",
"distillation": "distill",
"deliberation": "deliberation",
}
def _family_main(family):
if family == "sequential":
from outer.sequential import main
elif family == "hie":
from outer.mixture import main
elif family == "distill":
from outer.distillation import main
elif family == "deliberation":
from outer.deliberation import main
else: # pragma: no cover - guarded by argparse choices
raise ValueError(f"Unknown outer family: {family}")
return main
def main() -> None:
pre = argparse.ArgumentParser(add_help=False, description=__doc__)
pre.add_argument("--style", required=True, choices=sorted(STYLE_TO_FAMILY))
known, rest = pre.parse_known_args()
family = STYLE_TO_FAMILY[known.style]
family_main = _family_main(family)
try: # keep training output minimal: suppress datasets .map() progress bars
import datasets as _datasets
_datasets.disable_progress_bars()
except Exception:
pass
family_main(rest)
if __name__ == "__main__":
main()