56 lines
1.6 KiB
Python
56 lines
1.6 KiB
Python
import sys
|
|
from pathlib import Path
|
|
|
|
from litgpt.utils import CLI
|
|
|
|
|
|
def parser_commands() -> list[str]:
|
|
return [
|
|
"download",
|
|
"chat",
|
|
"finetune",
|
|
"finetune_lora",
|
|
"finetune_full",
|
|
"finetune_adapter",
|
|
"finetune_adapter_v2",
|
|
"pretrain",
|
|
"generate",
|
|
"generate_full",
|
|
"generate_adapter",
|
|
"generate_adapter_v2",
|
|
"generate_sequentially",
|
|
"generate_speculatively",
|
|
"generate_tp",
|
|
"convert_to_litgpt",
|
|
"convert_from_litgpt",
|
|
"convert_pretrained_checkpoint",
|
|
"merge_lora",
|
|
"evaluate",
|
|
"serve",
|
|
"validate",
|
|
]
|
|
|
|
|
|
def save_hyperparameters(
|
|
function: callable,
|
|
checkpoint_dir: Path,
|
|
known_commands: list[str] | None = None,
|
|
) -> None:
|
|
"""Captures the CLI parameters passed to `function` without running `function` and saves them to the checkpoint."""
|
|
from jsonargparse import capture_parser
|
|
|
|
# TODO: Make this more robust
|
|
# This hack strips away the subcommands from the top-level CLI
|
|
# to parse the file as if it was called as a script
|
|
if known_commands is None:
|
|
known_commands = parser_commands()
|
|
known_commands = [(c,) for c in known_commands]
|
|
for known_command in known_commands:
|
|
unwanted = slice(1, 1 + len(known_command))
|
|
if tuple(sys.argv[unwanted]) == known_command:
|
|
sys.argv[unwanted] = []
|
|
|
|
parser = capture_parser(lambda: CLI(function))
|
|
config = parser.parse_args()
|
|
parser.save(config, checkpoint_dir / "hyperparameters.yaml", overwrite=True)
|