From 135fc5741996431115128f15757748aed9efc315 Mon Sep 17 00:00:00 2001 From: Di Xu Date: Fri, 20 Mar 2026 10:03:30 -0700 Subject: [PATCH] Make extract_coreml_models.py buck-runnable Summary: Add a python_binary target for extract_coreml_models.py in the BUCK file, wrap the script's entry point in a main() function Differential Revision: D97498856 --- examples/apple/coreml/scripts/BUCK | 13 +++++++++++++ .../apple/coreml/scripts/extract_coreml_models.py | 11 ++++++++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/examples/apple/coreml/scripts/BUCK b/examples/apple/coreml/scripts/BUCK index 01f1d588ecd..164feb8d306 100644 --- a/examples/apple/coreml/scripts/BUCK +++ b/examples/apple/coreml/scripts/BUCK @@ -3,6 +3,19 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "no # targets.bzl. This file can contain fbcode-only targets. load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") +fbcode_target(_kind = python_binary, + name = "extract_coreml_models", + srcs = [ + "extract_coreml_models.py", + ], + main_function = "executorch.examples.apple.coreml.scripts.extract_coreml_models.main", + deps = [ + "//executorch/backends/apple/coreml:executorchcoreml", + "//executorch/exir:schema", + "//executorch/exir/_serialize:lib", + ], +) + fbcode_target(_kind = python_binary, name = "export", srcs = [ diff --git a/examples/apple/coreml/scripts/extract_coreml_models.py b/examples/apple/coreml/scripts/extract_coreml_models.py index dd72805d8c8..685b6b594f3 100644 --- a/examples/apple/coreml/scripts/extract_coreml_models.py +++ b/examples/apple/coreml/scripts/extract_coreml_models.py @@ -12,7 +12,6 @@ from typing import Dict, List, Optional from executorch.backends.apple.coreml import executorchcoreml -from executorch.backends.apple.coreml.compiler import CoreMLBackend from executorch.exir._serialize._program import deserialize_pte_binary from executorch.exir.schema import ( BackendDelegate, @@ -20,6 +19,8 @@ DataLocation, ) +COREML_BACKEND_ID = "CoreMLBackend" + def extract_coreml_models(pte_data: bytes): pte_file = deserialize_pte_binary(pte_data) @@ -39,7 +40,7 @@ def extract_coreml_models(pte_data: bytes): [execution_plan.delegates for execution_plan in program.execution_plan], [] ) coreml_delegates: List[BackendDelegate] = [ - delegate for delegate in delegates if delegate.id == CoreMLBackend.__name__ + delegate for delegate in delegates if delegate.id == COREML_BACKEND_ID ] # Track extracted models to avoid duplicates (multifunction models share partitions) @@ -109,7 +110,7 @@ def extract_coreml_models(pte_data: bytes): print("The model isn't delegated to Core ML.") -if __name__ == "__main__": +def main() -> None: """ Extracts the Core ML models embedded in the ``.pte`` file and saves them to the file system. @@ -127,3 +128,7 @@ def extract_coreml_models(pte_data: bytes): with open(model_path, mode="rb") as pte_file: pte_data = pte_file.read() extract_coreml_models(pte_data) + + +if __name__ == "__main__": + main()