File size: 2,957 Bytes
45cddd1 9dae136 4d9b54f 9dae136 45cddd1 8efff2e 9dae136 45cddd1 8efff2e 9dae136 8efff2e 45cddd1 8efff2e 9dae136 45cddd1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | use std::future::Future;
use std::path::{Path, PathBuf};
use anyhow::{Context, Result};
use candle_core::{DType, Device};
use candle_nn::VarBuilder;
use serde::de::DeserializeOwned;
pub fn model_dtype(device: &Device) -> DType {
crate::ops::model_dtype(device)
}
pub async fn resolve_manifest_path<F>(manifest: F) -> Result<PathBuf>
where
F: Future<Output = Result<PathBuf>>,
{
manifest.await
}
pub async fn load_mmaped_safetensors<F, T, Build, E>(
manifest: F,
device: &Device,
build: Build,
) -> Result<T>
where
F: Future<Output = Result<PathBuf>>,
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
let weights = resolve_manifest_path(manifest).await?;
load_mmaped_safetensors_path(&weights, device, build)
}
pub fn load_mmaped_safetensors_path<T, Build, E>(
weights: &Path,
device: &Device,
build: Build,
) -> Result<T>
where
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
load_mmaped_safetensors_path_with_dtype(weights, device, DType::F32, build)
}
pub fn load_mmaped_safetensors_path_with_dtype<T, Build, E>(
weights: &Path,
device: &Device,
dtype: DType,
build: Build,
) -> Result<T>
where
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[weights], dtype, device)? };
build(vb).map_err(Into::into)
}
pub async fn load_buffered_safetensors<F, T, Build, E>(
manifest: F,
device: &Device,
build: Build,
) -> Result<T>
where
F: Future<Output = Result<PathBuf>>,
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
let weights = resolve_manifest_path(manifest).await?;
load_buffered_safetensors_path(&weights, device, build)
}
pub fn load_buffered_safetensors_path<T, Build, E>(
weights: &Path,
device: &Device,
build: Build,
) -> Result<T>
where
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
load_buffered_safetensors_path_with_dtype(weights, device, DType::F32, build)
}
pub fn load_buffered_safetensors_path_with_dtype<T, Build, E>(
weights: &Path,
device: &Device,
dtype: DType,
build: Build,
) -> Result<T>
where
Build: FnOnce(VarBuilder) -> std::result::Result<T, E>,
E: Into<anyhow::Error>,
{
let data =
std::fs::read(weights).with_context(|| format!("failed to read {}", weights.display()))?;
let vb = VarBuilder::from_buffered_safetensors(data, dtype, device)?;
build(vb).map_err(Into::into)
}
pub fn read_json<T: DeserializeOwned>(path: &Path) -> Result<T> {
let data = std::fs::read_to_string(path)
.with_context(|| format!("failed to read {}", path.display()))?;
let parsed = serde_json::from_str(&data)
.with_context(|| format!("failed to parse {}", path.display()))?;
Ok(parsed)
}
|