Function create_new_metadata

cake-core/src/utils/split.rs:115–149

(
    data_path: &Path,
    reduced: &HashMap<String, Vec<String>>,
)

}

fn create_new_metadata(
    data_path: &Path,
    reduced: &HashMap<String, Vec<String>>,
) -> Result<HashMap<String, TensorStore>> {
    let mut metadata: HashMap<String, TensorStore> = HashMap::new();

    for (filename, tensor_names) in reduced {
        let filepath = data_path.join(filename);

        log::info!("loading {} ...", filepath.display());

        let file = File::open(&filepath)?;
        let buffer = unsafe { memmap2::MmapOptions::new().map(&file)? };
        let tensors = SafeTensors::deserialize(&buffer)?;

        log::info!("  extracting {} tensors", tensor_names.len());

        for tensor_name in tensor_names {
            let tensor = tensors.tensor(tensor_name)?;
            metadata.insert(
                tensor_name.to_string(),
                TensorStore {
                    dtype: tensor.dtype(),
                    shape: tensor.shape().to_vec(),
                    data: tensor.data().to_vec(),
                },
            );
        }

        drop(tensors);
        drop(buffer);
    }

    Ok(metadata)
}

/// Split a model into per-worker bundles.
///