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 simple models to flatbuffer 8 9import logging 10 11from executorch.backends.cadence.aot.ops_registrations import * # noqa 12 13import torch 14 15from executorch.backends.cadence.aot.export_example import export_model 16 17from executorch.examples.models.llama.llama_transformer import ModelArgs, Transformer 18 19 20FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" 21logging.basicConfig(level=logging.INFO, format=FORMAT) 22 23 24def main() -> None: 25 args = ModelArgs( 26 dim=512, 27 vocab_size=512, 28 hidden_dim=1024, 29 n_heads=8, 30 # use_kv_cache=True, 31 n_layers=1, 32 ) 33 seq = 64 34 b = 1 35 model = Transformer(args) 36 example_inputs = (torch.randint(0, 10, [b, seq], dtype=torch.int64),) 37 38 export_model(model, example_inputs) 39 40 41if __name__ == "__main__": 42 main() 43