Load a model of any supported format.
(path: Path)
| 1089 | |
| 1090 | |
| 1091 | def load_some_model(path: Path) -> ModelPlus: |
| 1092 | '''Load a model of any supported format.''' |
| 1093 | # Be extra-friendly and accept either a file or a directory: |
| 1094 | if path.is_dir(): |
| 1095 | # Check if it's a set of safetensors files first |
| 1096 | globs = ["model-00001-of-*.safetensors", "model.safetensors"] |
| 1097 | files = [file for glob in globs for file in path.glob(glob)] |
| 1098 | if not files: |
| 1099 | # Try the PyTorch patterns too, with lower priority |
| 1100 | globs = ["consolidated.00.pth", "pytorch_model-00001-of-*.bin", "*.pt", "pytorch_model.bin"] |
| 1101 | files = [file for glob in globs for file in path.glob(glob)] |
| 1102 | if not files: |
| 1103 | raise Exception(f"Can't find model in directory {path}") |
| 1104 | if len(files) > 1: |
| 1105 | raise Exception(f"Found multiple models in {path}, not sure which to pick: {files}") |
| 1106 | path = files[0] |
| 1107 | |
| 1108 | paths = find_multifile_paths(path) |
| 1109 | models_plus: list[ModelPlus] = [] |
| 1110 | for path in paths: |
| 1111 | print(f"Loading model file {path}") |
| 1112 | models_plus.append(lazy_load_file(path)) |
| 1113 | |
| 1114 | model_plus = merge_multifile_models(models_plus) |
| 1115 | return model_plus |
| 1116 | |
| 1117 | def load_predictor_model(path: Path) -> ModelPlus: |
| 1118 | '''Load MLP models for sparse FFN inference from directory.''' |
no test coverage detected