| 43 | |
| 44 | |
| 45 | def load_source_index(model_dir: Path) -> dict: |
| 46 | index_path = model_dir / "model.safetensors.index.json" |
| 47 | if index_path.exists(): |
| 48 | with open(index_path) as f: |
| 49 | return json.load(f) |
| 50 | single = model_dir / "model.safetensors" |
| 51 | if single.exists(): |
| 52 | with safe_open(str(single), framework="pt") as f: |
| 53 | return { |
| 54 | "metadata": {}, |
| 55 | "weight_map": {k: "model.safetensors" for k in f.keys()}, |
| 56 | } |
| 57 | raise FileNotFoundError(f"No safetensors files found in {model_dir}") |
| 58 | |
| 59 | |
| 60 | def load_source_config(model_dir: Path) -> dict: |