nano_gpt.tool.export

 1"""Command-line interface for exporting a model checkpoint to safet tensor format.
 2
 3Usage:
 4```
 5usage: nano-gpt export [-h] [--pretrained {gpt2,gpt2-large,gpt2-medium,gpt2-xl}]
 6                       [--model {gpt2,gpt2-large,gpt2-medium,gpt2-xl,gpt2-xs,gpt2-xxs}] [--checkpoint CHECKPOINT] [--device DEVICE]
 7                       [--sequence-length SEQUENCE_LENGTH] [--seed SEED] [--compile | --no-compile] --export-path EXPORT_PATH
 8
 9Export a model checkpoint to safetensors
10
11options:
12  -h, --help            show this help message and exit
13
14model:
15  --pretrained {gpt2,gpt2-large,gpt2-medium,gpt2-xl}
16                        The name of the pretrained model to use.
17  --model {gpt2,gpt2-large,gpt2-medium,gpt2-xl,gpt2-xs,gpt2-xxs}
18                        Use the specified model name configuration default values.
19  --checkpoint CHECKPOINT
20                        Load a model from a checkpoint.
21  --device DEVICE       The device to use.
22  --sequence-length SEQUENCE_LENGTH
23                        The sequence length used for input content in each micro batch.
24  --seed SEED           The seed to use for sampling/training.
25  --compile, --no-compile
26                        Will compile the model if supported by the device.
27
28export:
29  --export-path EXPORT_PATH
30                        Destination model path for exporting a checkpoint to safetensors.
31```
32"""
33
34import argparse
35import logging
36import pathlib
37
38from nano_gpt.checkpoint import export_checkpoint
39
40from .model_config import (
41    create_model_arguments,
42    load_checkpoint_context,
43)
44
45_LOGGER = logging.getLogger(__name__)
46
47
48def create_arguments(args: argparse.ArgumentParser) -> None:
49    """Get parsed passed in arguments."""
50    create_model_arguments(args)
51    group = args.add_argument_group("export")
52    group.add_argument(
53        "--export-dir",
54        type=str,
55        help="Destination directory path for exporting a checkpoint to safetensors.",
56        required=True,
57    )
58
59
60def run(args: argparse.Namespace) -> int:
61    """Run the export command."""
62
63    with load_checkpoint_context(args) as checkpoint:
64        if checkpoint is None:
65            raise ValueError("Required flag --checkpoint not specified or not found")
66        checkpoint_path = pathlib.Path(args.checkpoint)
67        export_path = pathlib.Path(args.export_dir)
68        export_checkpoint(
69            checkpoint,
70            checkpoint_path,
71            export_path,
72        )
73
74    return 0