Function load_var_builder_from_index

cake-core/src/utils/mod.rs:251–267

Create a VarBuilder with the tensors loaded from the index.

(
    tensor_index: PathBuf,
    dtype: DType,
    device: Device,
    quant: &dyn Quantization,
)

/// Create a VarBuilder with the tensors loaded from the index.
pub fn load_var_builder_from_index<'a>(
    tensor_index: PathBuf,
    dtype: DType,
    device: Device,
    quant: &dyn Quantization,
) -> Result<VarBuilder<'a>> {
    let filenames: Vec<std::path::PathBuf> = if tensor_index.exists() {
        load_safetensors_paths_from_index(tensor_index)
            .map_err(|e| anyhow!("can't load tensors index: {:?}", e))?
    } else {
        load_safetensors_from_model(tensor_index.parent().unwrap())
            .map_err(|e| anyhow!("can't load tensors index: {:?}", e))?
    };

    prefetch_safetensors(&filenames)?;
    unsafe { quant.load_var_builder(&filenames, dtype, &device) }
}

/// Create a VarBuilder that only loads safetensors shards needed for the given
/// local layers. Shards containing only remote-worker tensors are excluded,