| 86 | } |
| 87 | |
| 88 | fn reduce_for_worker( |
| 89 | index: &Index, |
| 90 | worker: &Node, |
| 91 | ) -> Result<(Index, HashMap<String, Vec<String>>)> { |
| 92 | log::info!("worker: {}", &worker.host); |
| 93 | |
| 94 | let mut reduced: HashMap<String, Vec<String>> = HashMap::new(); |
| 95 | let mut new_index = Index::new(); |
| 96 | |
| 97 | for (layer_full_name, filename) in &index.weight_map { |
| 98 | if worker.is_text_model_layer_owner(layer_full_name) { |
| 99 | if let Some(layers) = reduced.get_mut(filename) { |
| 100 | layers.push(layer_full_name.to_string()); |
| 101 | } else { |
| 102 | reduced.insert(filename.to_string(), vec![layer_full_name.to_string()]); |
| 103 | } |
| 104 | |
| 105 | new_index.weight_map.insert( |
| 106 | layer_full_name.to_string(), |
| 107 | "reduced.safetensors".to_string(), |
| 108 | ); |
| 109 | } |
| 110 | } |
| 111 | |
| 112 | Ok((new_index, reduced)) |
| 113 | } |
| 114 | |
| 115 | fn create_new_metadata( |
| 116 | data_path: &Path, |