Build from a multi-shard model with index JSON.
(index_path: &Path)
| 308 | |
| 309 | /// Build from a multi-shard model with index JSON. |
| 310 | fn from_index_json(index_path: &Path) -> Result<Self> { |
| 311 | let parent = index_path |
| 312 | .parent() |
| 313 | .ok_or_else(|| anyhow::anyhow!("no parent dir"))?; |
| 314 | |
| 315 | // Parse the weight_map: tensor_name → shard_filename |
| 316 | let json: serde_json::Value = |
| 317 | serde_json::from_reader(File::open(index_path)?)?; |
| 318 | let weight_map = json |
| 319 | .get("weight_map") |
| 320 | .and_then(|v| v.as_object()) |
| 321 | .ok_or_else(|| anyhow::anyhow!("missing weight_map in index"))?; |
| 322 | |
| 323 | // Collect unique shard filenames (we get tensor metadata from shard headers, not weight_map) |
| 324 | let mut shard_names: Vec<&str> = weight_map.values() |
| 325 | .filter_map(|v| v.as_str()) |
| 326 | .collect(); |
| 327 | shard_names.sort_unstable(); |
| 328 | shard_names.dedup(); |
| 329 | |
| 330 | let mut index = HashMap::with_capacity(weight_map.len()); |
| 331 | let mut shards = Vec::with_capacity(shard_names.len()); |
| 332 | |
| 333 | for shard_name in shard_names { |
| 334 | let shard_path = parent.join(shard_name); |
| 335 | let shard_idx = shards.len() as u16; |
| 336 | let f = File::open(&shard_path)?; |
| 337 | #[cfg(unix)] |
| 338 | let shard = MappedShard::new(&f)?; |
| 339 | #[cfg(unix)] |
| 340 | let (header_len, header) = Self::parse_shard_header_mmap(&shard)?; |
| 341 | #[cfg(not(unix))] |
| 342 | let (header_len, header) = Self::parse_shard_header(&shard_path)?; |
| 343 | #[cfg(not(unix))] |
| 344 | let shard = MappedShard::new(f)?; |
| 345 | shards.push(shard); |
| 346 | |
| 347 | for (name, entry) in header { |
| 348 | if name.starts_with("__") { |
| 349 | continue; |
| 350 | } |
| 351 | if let Some(tm) = Self::parse_tensor_meta(&name, &entry, shard_idx, header_len)? { |
| 352 | index.insert(name, tm); |
| 353 | } |
| 354 | } |
| 355 | } |
| 356 | |
| 357 | log::info!( |
| 358 | "SafetensorsStorage: indexed {} tensors across {} shards", |
| 359 | index.len(), |
| 360 | shards.len() |
| 361 | ); |
| 362 | |
| 363 | Ok(Self { index, shards }) |
| 364 | } |
| 365 | |
| 366 | /// Build from a single safetensors file (non-sharded model). |
| 367 | fn from_single_file(path: &Path) -> Result<Self> { |