Convert model checkpoints to huggingface format.
(
checkpoint_dir: Annotated[
str,
typer.Option("--checkpoint-dir", "-c", help="The path to the checkpoint directory."),
],
base_model_dir: Annotated[
Optional[str],
typer.Option("--base-model-dir", "-b", help="The path to the base model."),
] = None,
step: Annotated[
Optional[str],
typer.Option(
"--step",
"-s",
help="Specific step number(s) to convert. Comma-separated (e.g., 100,200,300) or repeated (-s 100 -s 200).",
),
] = None,
)
| 6 | |
| 7 | |
| 8 | def convert_command( |
| 9 | checkpoint_dir: Annotated[ |
| 10 | str, |
| 11 | typer.Option("--checkpoint-dir", "-c", help="The path to the checkpoint directory."), |
| 12 | ], |
| 13 | base_model_dir: Annotated[ |
| 14 | Optional[str], |
| 15 | typer.Option("--base-model-dir", "-b", help="The path to the base model."), |
| 16 | ] = None, |
| 17 | step: Annotated[ |
| 18 | Optional[str], |
| 19 | typer.Option( |
| 20 | "--step", |
| 21 | "-s", |
| 22 | help="Specific step number(s) to convert. Comma-separated (e.g., 100,200,300) or repeated (-s 100 -s 200).", |
| 23 | ), |
| 24 | ] = None, |
| 25 | ) -> None: |
| 26 | """Convert model checkpoints to huggingface format.""" |
| 27 | from trinity.manager.checkpoint_converter import Converter |
| 28 | |
| 29 | converter = Converter(base_model_dir) |
| 30 | |
| 31 | # Parse step parameter (supports both "100,200,300" and multiple -s flags) |
| 32 | step_list: List[int] = [] |
| 33 | if step: |
| 34 | # Split by comma and/or whitespace, then convert to int |
| 35 | for part in step.replace(",", " ").split(): |
| 36 | try: |
| 37 | step_list.append(int(part)) |
| 38 | except ValueError: |
| 39 | typer.echo(f"[ERROR] Invalid step number: {part}", err=True) |
| 40 | raise typer.Exit(code=1) |
| 41 | |
| 42 | if step_list: |
| 43 | _convert_multi_steps(step_list, checkpoint_dir, converter) |
| 44 | else: |
| 45 | # Original behavior: convert all or a single checkpoint |
| 46 | dir_path = checkpoint_dir |
| 47 | if "global_step_" in dir_path: |
| 48 | while not os.path.basename(dir_path).startswith("global_step_"): |
| 49 | dir_path = os.path.dirname(dir_path) |
| 50 | converter.convert(dir_path) |
| 51 | |
| 52 | |
| 53 | def _convert_multi_steps(step_list: List[int], checkpoint_dir: str, converter) -> None: |