(
f: Union[str, os.PathLike[str], io.BytesIO],
*,
extra_files: Optional[Dict[str, Any]] = None,
)
| 770 | |
| 771 | |
| 772 | def load( |
| 773 | f: Union[str, os.PathLike[str], io.BytesIO], |
| 774 | *, |
| 775 | extra_files: Optional[Dict[str, Any]] = None, |
| 776 | ) -> ep.ExportedProgram: |
| 777 | if isinstance(f, (str, os.PathLike)): |
| 778 | f = os.fspath(str(f)) |
| 779 | |
| 780 | extra_files = extra_files or {} |
| 781 | |
| 782 | with zipfile.ZipFile(f, "r") as zipf: |
| 783 | # Check the version |
| 784 | version = zipf.read("version").decode().split(".") |
| 785 | |
| 786 | assert len(version) == len(SCHEMA_VERSION) |
| 787 | if version[0] != str(SCHEMA_VERSION[0]): |
| 788 | raise RuntimeError( |
| 789 | f"Serialized version {version} does not match our current " |
| 790 | f"schema version {SCHEMA_VERSION}." |
| 791 | ) |
| 792 | |
| 793 | # Load serialized_ep and serialized_state_dict from the zip file |
| 794 | |
| 795 | serialized_exported_program: Optional[bytes] = None |
| 796 | serialized_state_dict: Optional[bytes] = None |
| 797 | serialized_constants: Optional[bytes] = None |
| 798 | serialized_example_inputs: Optional[bytes] = None |
| 799 | |
| 800 | for file_info in zipf.infolist(): |
| 801 | file_content = zipf.read(file_info.filename) |
| 802 | |
| 803 | if file_info.filename == "serialized_exported_program.json": |
| 804 | serialized_exported_program = file_content |
| 805 | elif file_info.filename == "serialized_state_dict.json": |
| 806 | print("This version of file is deprecated") |
| 807 | serialized_state_dict = file_content |
| 808 | elif file_info.filename == "serialized_constants.json": |
| 809 | print("This version of file is deprecated") |
| 810 | serialized_constants = file_content |
| 811 | elif file_info.filename == "serialized_state_dict.pt": |
| 812 | serialized_state_dict = file_content |
| 813 | elif file_info.filename == "serialized_constants.pt": |
| 814 | serialized_constants = file_content |
| 815 | elif file_info.filename.startswith("extra_files"): |
| 816 | filename = file_info.filename.split("/", 1)[1] |
| 817 | extra_files[filename] = file_content.decode("utf-8") |
| 818 | elif file_info.filename == "serialized_example_inputs.pt": |
| 819 | serialized_example_inputs = file_content |
| 820 | |
| 821 | assert serialized_exported_program is not None |
| 822 | assert serialized_state_dict is not None |
| 823 | assert serialized_constants is not None |
| 824 | assert serialized_example_inputs is not None |
| 825 | |
| 826 | artifact: export_serialize.SerializedArtifact = ( |
| 827 | export_serialize.SerializedArtifact( |
| 828 | serialized_exported_program, |
| 829 | serialized_state_dict, |
no test coverage detected