Make export file compatible with buck (#5703)

Summary:
Pull Request resolved: #5703

Part of the export.py isn't working well with buck, move the function from `__main__.py` to a standalone `main` function so it can be called
ghstack-source-id: 244927520
exported-using-ghexport
bypass-github-export-checks

Reviewed By: shoumikhin

Differential Revision: D63476613

fbshipit-source-id: 1defdde080f05b5c31557de54314f3423b7e3c4f

Author: cccclai

examples/apple/coreml/scripts/export.py

@@ -21,14 +21,15 @@
 from executorch.exir import to_edge
 
 from executorch.exir.backend.backend_api import to_backend
+
 from torch.export import export
 
 REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent.parent.parent.parent
 EXAMPLES_DIR = REPO_ROOT / "examples"
 sys.path.append(str(EXAMPLES_DIR.absolute()))
 
-from models import MODEL_NAME_TO_MODEL
-from models.model_factory import EagerModelFactory
+from executorch.examples.models import MODEL_NAME_TO_MODEL
+from executorch.examples.models.model_factory import EagerModelFactory
 
 # Script to export a model with coreml delegation.
 
@@ -82,7 +83,7 @@ def partition_module_to_coreml(module):
     module = module.eval()
 
 
-def lower_module_to_coreml(module, compile_specs):
+def lower_module_to_coreml(module, compile_specs, example_inputs):
     module = module.eval()
     edge = to_edge(export(module, example_inputs), compile_config=_EDGE_COMPILE_CONFIG)
     # All of the subsequent calls on the edge_dialect_graph generated above (such as delegation or
@@ -141,7 +142,7 @@ def generate_compile_specs_from_args(args):
     )
 
 
-if __name__ == "__main__":
+def main():
     args = parse_args()
 
     if args.model_name not in MODEL_NAME_TO_MODEL:
@@ -167,6 +168,7 @@ def generate_compile_specs_from_args(args):
     if args.use_partitioner:
         model.eval()
         exir_program_aten = torch.export.export(model, example_inputs)
+
         edge_program_manager = exir.to_edge(exir_program_aten)
         edge_copy = copy.deepcopy(edge_program_manager)
         partitioner = CoreMLPartitioner(
@@ -179,6 +181,7 @@ def generate_compile_specs_from_args(args):
     else:
         lowered_module, edge_copy = lower_module_to_coreml(
             module=model,
+            example_inputs=example_inputs,
             compile_specs=compile_specs,
         )
         exec_program = export_lowered_module_to_executorch_program(
@@ -193,3 +196,7 @@ def generate_compile_specs_from_args(args):
         save_processed_bytes(
             lowered_module.processed_bytes, args.model_name, args.compute_unit
         )
+
+
+if __name__ == "__main__":
+    main()