Spaces:
Sleeping
Sleeping
File size: 2,573 Bytes
be9f416 92e9724 be9f416 92e9724 be9f416 92e9724 be9f416 92e9724 be9f416 92e9724 be9f416 | 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 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """Group reward optimization entrypoint.
Dry-run writes a reproducible manifest. Full mode runs a lightweight
group-relative policy-gradient loop with Hugging Face Transformers/PyTorch.
"""
from __future__ import annotations
import argparse
try:
from ontology_architect.config import load_config
from ontology_architect.training import run_group_reward_optimization, write_gro_manifest
except ImportError: # pragma: no cover
from config import load_config
from training import run_group_reward_optimization, write_gro_manifest
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", default="configs/tiny_smoke.json")
parser.add_argument("--model-id", required=True)
parser.add_argument("--output-dir", default=None)
parser.add_argument("--group-size", type=int, default=None)
parser.add_argument("--max-steps", type=int, default=None)
parser.add_argument("--max-new-tokens", type=int, default=768)
parser.add_argument("--learning-rate", type=float, default=1e-6)
parser.add_argument("--checkpoint-steps", type=int, default=None)
parser.add_argument("--dry-run", action="store_true")
args = parser.parse_args()
config = load_config(args.config)
output_dir = args.output_dir or f"{config.training.output_dir}/gro"
group_size = args.group_size or config.training.group_size
max_steps = args.max_steps or config.training.max_steps
checkpoint_steps = args.checkpoint_steps if args.checkpoint_steps is not None else config.training.checkpoint_steps
if args.dry_run:
manifest = write_gro_manifest(
args.model_id,
output_dir,
group_size,
max_steps,
dry_run=True,
checkpoint_steps=checkpoint_steps,
)
print(
f"validated GRO config: group_size={manifest['group_size']} "
f"max_steps={manifest['max_steps']} checkpoint_steps={manifest['checkpoint_steps']}"
)
return
result = run_group_reward_optimization(
args.model_id,
config,
output_dir,
group_size,
max_steps,
max_new_tokens=args.max_new_tokens,
learning_rate=args.learning_rate,
checkpoint_steps=checkpoint_steps,
)
write_gro_manifest(
args.model_id,
output_dir,
group_size,
max_steps,
dry_run=False,
checkpoint_steps=checkpoint_steps,
)
print(f"saved GRO checkpoint to {result['output_dir']}")
if __name__ == "__main__":
main()
|