mirror of
https://wget.la/https://github.com/leookun/cursor-byok
synced 2026-10-06 05:12:03 +08:00
88 lines
3.0 KiB
Rust
88 lines
3.0 KiB
Rust
//! Verified, atomic download and reuse of the tokenizer and embedding matrix.
|
|
|
|
use std::{
|
|
fs,
|
|
io::{Read, Write},
|
|
path::{Path, PathBuf},
|
|
};
|
|
|
|
use sha2::{Digest, Sha256};
|
|
|
|
use crate::{Error, Result};
|
|
|
|
const MODEL_URL: &str =
|
|
"https://huggingface.co/minishlab/potion-code-16M-v2/resolve/main/model.safetensors";
|
|
const TOKENIZER_URL: &str =
|
|
"https://huggingface.co/minishlab/potion-code-16M-v2/resolve/main/tokenizer.json";
|
|
const MODEL_SHA256: &str = "75cf7a6c2171b230ad19b1e7d8e0b1aee86da5a02af8e7cacedd9921d227623c";
|
|
const TOKENIZER_SHA256: &str = "107bbdcbad4bff1d299b7a4c3a2fb17c52890688b7dd0e4c9deab79d3c4f3d45";
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct ModelAssets {
|
|
pub model: PathBuf,
|
|
pub tokenizer: PathBuf,
|
|
}
|
|
|
|
impl ModelAssets {
|
|
pub fn model_path(cache_root: &Path) -> PathBuf {
|
|
cache_root.join("models/potion-code-16M-v2/model.safetensors")
|
|
}
|
|
|
|
pub fn ensure(cache_root: &Path) -> Result<Self> {
|
|
let directory = cache_root.join("models/potion-code-16M-v2");
|
|
fs::create_dir_all(&directory).map_err(|error| Error::io(&directory, error))?;
|
|
let model = Self::model_path(cache_root);
|
|
let tokenizer = directory.join("tokenizer.json");
|
|
ensure_asset(&model, MODEL_URL, MODEL_SHA256)?;
|
|
ensure_asset(&tokenizer, TOKENIZER_URL, TOKENIZER_SHA256)?;
|
|
Ok(Self { model, tokenizer })
|
|
}
|
|
}
|
|
|
|
fn ensure_asset(path: &Path, url: &str, expected: &str) -> Result<()> {
|
|
if path.is_file() && digest(path)? == expected {
|
|
return Ok(());
|
|
}
|
|
let temporary = path.with_extension(format!("tmp-{}", std::process::id()));
|
|
let response = reqwest::blocking::Client::builder()
|
|
.build()
|
|
.map_err(|error| Error::ModelAsset(error.to_string()))?
|
|
.get(url)
|
|
.send()
|
|
.and_then(reqwest::blocking::Response::error_for_status)
|
|
.map_err(|error| Error::ModelAsset(format!("download {url}: {error}")))?;
|
|
let bytes = response
|
|
.bytes()
|
|
.map_err(|error| Error::ModelAsset(error.to_string()))?;
|
|
let actual = hex::encode(Sha256::digest(&bytes));
|
|
if actual != expected {
|
|
return Err(Error::ModelAsset(format!(
|
|
"checksum mismatch for {url}: expected {expected}, got {actual}"
|
|
)));
|
|
}
|
|
let mut file = fs::File::create(&temporary).map_err(|error| Error::io(&temporary, error))?;
|
|
file.write_all(&bytes)
|
|
.and_then(|_| file.sync_all())
|
|
.map_err(|error| Error::io(&temporary, error))?;
|
|
if path.exists() {
|
|
fs::remove_file(path).map_err(|error| Error::io(path, error))?;
|
|
}
|
|
fs::rename(&temporary, path).map_err(|error| Error::io(path, error))
|
|
}
|
|
|
|
fn digest(path: &Path) -> Result<String> {
|
|
let mut file = fs::File::open(path).map_err(|error| Error::io(path, error))?;
|
|
let mut hash = Sha256::new();
|
|
let mut buffer = [0_u8; 64 * 1024];
|
|
loop {
|
|
let read = file
|
|
.read(&mut buffer)
|
|
.map_err(|error| Error::io(path, error))?;
|
|
if read == 0 {
|
|
break;
|
|
}
|
|
hash.update(&buffer[..read]);
|
|
}
|
|
Ok(hex::encode(hash.finalize()))
|
|
}
|