xref: /aosp_15_r20/external/executorch/examples/models/llama/export_llama.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a)
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7# Example script for exporting Llama2 to flatbuffer
8
9import logging
10import sys
11
12import torch
13
14from .export_llama_lib import build_args_parser, export_llama
15
16sys.setrecursionlimit(4096)
17
18
19FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
20logging.basicConfig(level=logging.INFO, format=FORMAT)
21
22
23def main() -> None:
24    seed = 42
25    torch.manual_seed(seed)
26    parser = build_args_parser()
27    args = parser.parse_args()
28    export_llama(args)
29
30
31if __name__ == "__main__":
32    main()  # pragma: no cover
33