(data_path: &Path)
| 53 | } |
| 54 | |
| 55 | fn load_index(data_path: &Path) -> Result<Index> { |
| 56 | let tensors_index_path = data_path.join("model.safetensors.index.json"); |
| 57 | |
| 58 | if tensors_index_path.exists() { |
| 59 | let tensors_index_data = std::fs::read_to_string(tensors_index_path)?; |
| 60 | let tensors_index: Index = serde_json::from_str(&tensors_index_data)?; |
| 61 | Ok(tensors_index) |
| 62 | } else { |
| 63 | let single_path = data_path.join("model.safetensors"); |
| 64 | if !single_path.exists() { |
| 65 | anyhow::bail!( |
| 66 | "neither model.safetensors.index.json nor model.safetensors found in {}", |
| 67 | data_path.display() |
| 68 | ); |
| 69 | } |
| 70 | |
| 71 | log::info!("no index file found, generating from model.safetensors ..."); |
| 72 | |
| 73 | let file = File::open(&single_path)?; |
| 74 | let buffer = unsafe { memmap2::MmapOptions::new().map(&file)? }; |
| 75 | let tensors = SafeTensors::deserialize(&buffer)?; |
| 76 | |
| 77 | let mut index = Index::new(); |
| 78 | for (name, _) in tensors.tensors() { |
| 79 | index |
| 80 | .weight_map |
| 81 | .insert(name.to_string(), "model.safetensors".to_string()); |
| 82 | } |
| 83 | |
| 84 | Ok(index) |
| 85 | } |
| 86 | } |
| 87 | |
| 88 | fn reduce_for_worker( |
| 89 | index: &Index, |
no outgoing calls