[dependencies] # tokenizers 是 HuggingFace 官方 Rust 实现，Python 的 transformers 底层就是它 tokenizers = { version = "0.21", features = ["http"] } # 生产上强烈建议把 tokenizer.json 一起打进镜像/包，不要运行时去 HF 下载 // src/tokenize.rs use tokenizers::Tokenizer; use std::path::Path; pub fn load_tokenizer(path: &Path) -> anyhow::Result<Tokenizer> { // from_file 会自动识别 tokenizer.json（HF 的"自包含"格式，含词表+规则+后处理） let tokenizer = Tokenizer::from_file(path) .map_err(|e| anyhow::anyhow!("加载 tokenizer 失败 {}: {e}", path.display()))?; // 加载后必须做的自检：模型真正需要的那几个 special token 在不在 // 缺 CLS/SEP/PAD 时不会报错，只会让结果悄悄变差 —— 这是最恶心的坑 for name in ["[CLS]", "[SEP]", "[PAD]", "[UNK]"] { if tokenizer.token_to_id(name).is_none() { anyhow::bail!("tokenizer 缺少 special token {name}，检查是否用了配套的词表"); } } Ok(tokenizer) } // 顺手把关键信息打出来，部署时对不上版本第一时间就能发现 pub fn describe(t: &Tokenizer) { println!("vocab_size = {}", t.get_vocab_size(true)); println!("cls = {:?}, sep = {:?}, pad = {:?}", t.token_to_id("[CLS]"), t.token_to_id("[SEP]"), t.token_to_id("[PAD]")); // 一个非常实用的"指纹"：对固定文本分词，把 id 序列打印出来 // 换模型/换分词器时对比这一串，能立刻发现不匹配 if let Ok(e) = t.encode("版本自检", true) { println!("probe ids = {:?}", e.get_ids()); } }

use tokenizers::Tokenizer; let tokenizer: Tokenizer = load_tokenizer(Path::new("models/bert/tokenizer.json"))?; // 第二个参数 add_special_tokens = true 会按后处理模板自动加 [CLS] / [SEP] // 千万不要自己手动拼 special token，交给 tokenizer 的 post_processor 做 let enc = tokenizer.encode("这家餐厅的服务态度非常好", true)?; // 1) input_ids：送进模型的主输入 let ids: &[u32] = enc.get_ids(); // 2) attention_mask：1 = 真实 token，0 = padding。缺了它模型会把 pad 当内容算 let mask: &[u32] = enc.get_attention_mask(); // 3) token_type_ids：用于"句子对"任务区分第一句/第二句；单句任务全 0 // BERT 类模型需要它，RoBERTa / 很多新模型不需要 let type_ids: &[u32] = enc.get_type_ids(); // 4) tokens：还原成人能看懂的 token 串，调试分词结果时必看 println!("tokens = {:?}", enc.get_tokens()); // 5) special_tokens_mask：标出哪些位置是 [CLS]/[SEP]/[PAD]，后处理时用得到 let special = enc.get_special_tokens_mask(); // 6) offsets：每个 token 对应原文的字节区间，做实体抽取(NER)时必需 let offsets = enc.get_offsets(); // decode：把 id 还原成文本。skip_special_tokens 一般要开，否则会带出 [CLS] let text = tokenizer.decode(ids, true)?; println!("还原 = {text}"); // 一段可直接放进单元测试的自检代码 #[test] fn tokenizer_fingerprint_is_stable() { let t = load_tokenizer(Path::new("models/bert/tokenizer.json")).unwrap(); let e = t.encode("版本自检", true).unwrap(); // 换 tokenizer 文件时这个断言会失败，逼你确认"是故意换的还是搞错了" assert_eq!(e.get_ids(), &[101, 2345, 6789, 102]); }

use tokenizers::{PaddingParams, PaddingStrategy, TruncationParams, TruncationStrategy, TruncationDirection, Tokenizer}; pub fn configure(tokenizer: &mut Tokenizer, max_len: usize) -> anyhow::Result<()> { // 1) padding 策略：批内最长（动态，省算力）还是固定长度（形状恒定，省重编译） tokenizer.with_padding(Some(PaddingParams { // BatchLongest：补齐到这一批里最长的长度 —— 推荐的默认值 // Fixed(len)：永远补齐到固定长度 —— 静态 shape 模型必须用它 strategy: PaddingStrategy::BatchLongest, // 右侧 padding 是 BERT 类模型的默认约定；GPT 类是左侧 direction: tokenizers::PaddingDirection::Right, // 补齐到 8 的整数倍，便于某些硬件/算子做向量化 pad_to_multiple_of: Some(8), // pad_token / pad_id 一般不用改，tokenizer.json 里已经声明好了 ..Default::default() })); // 2) truncation 策略：超长时从哪里砍 tokenizer.with_truncation(Some(TruncationParams { max_length: max_len, // LongestFirst：先砍最长的序列（批处理时能让长度分布更均匀） strategy: TruncationStrategy::LongestFirst, // stride = 0 表示不重叠；做滑窗长文本时要设成 > 0 stride: 0, // 保留开头还是结尾：新闻/长文一般保留开头（Direction::Right） // 但"关键词在结尾"的场景（比如客服对话）就该用 Left direction: TruncationDirection::Right, }))?; Ok(()) } // 配好之后每次 encode 自动生效，不用在调用处重复传参 let enc = tokenizer.encode("很长的文本...", true)?; assert!(enc.get_ids().len() <= max_len);

const MAX_LEN: usize = 512; pub struct Prepared { pub ids: Vec<i64>, pub mask: Vec<i64>, // 把这两个数带出去，指标里就能看到"有多少请求被截断" pub raw_len: usize, pub truncated: bool, } pub fn prepare(tokenizer: &Tokenizer, text: &str) -> anyhow::Result<Prepared> { // 先不加 truncation 编一次，拿到"真实长度"用于判断 // 大文本场景下这次额外的编码开销可接受；追求极限性能可以改用 tokenizer.tokenize 估长 let full = tokenizer.encode(text, true)?; let raw_len = full.get_ids().len(); let (ids, mask, truncated) = if raw_len > MAX_LEN { // 走配置好的 truncation 路径 let enc = tokenizer.encode(text, true)?; (enc.get_ids(), enc.get_attention_mask(), true) } else { (full.get_ids(), full.get_attention_mask(), false) }; if truncated { // 关键：截断一定要有日志/指标，否则就是静默丢数据 tracing::warn!( raw_tokens = raw_len, kept_tokens = ids.len(), "输入超过模型上限，已截断" ); } Ok(Prepared { ids: ids.iter().map(|&i| i as i64).collect(), mask: mask.iter().map(|&m| m as i64).collect(), raw_len, truncated, }) }

use std::sync::Arc; use axum::{extract::State, http::StatusCode, routing::post, Json, Router}; use serde::{Deserialize, Serialize}; use tokio::sync::{mpsc, oneshot}; use tokenizers::Tokenizer; #[derive(Deserialize)] pub struct PredictReq { pub text: String } #[derive(Serialize)] pub struct PredictResp { pub label: String, pub score: f32, // 把可观测性信息返回给调用方，排查时非常有用 pub tokens: usize, pub truncated: bool, pub latency_ms: u64, } pub struct ModelWorker { // tokenizer 是 Send + Sync 的，可以直接 Arc 共享给多个请求处理器 pub tokenizer: Arc<Tokenizer>, // 一个 channel 就把"并发请求"收敛成"串行推理" pub tx: mpsc::Sender<Job>, // 用 Semaphore 做背压：队列满了要快速失败，而不是无限堆积 pub inflight: Arc<tokio::sync::Semaphore>, } pub struct Job { pub ids: Vec<i64>, pub mask: Vec<i64>, pub reply: oneshot::Sender<anyhow::Result<(usize, f32)>>, } pub async fn predict( State(w): State<Arc<ModelWorker>>, Json(req): Json<PredictReq>, ) -> Result<Json<PredictResp>, (StatusCode, String)> { let started = std::time::Instant::now(); // 1) 背压：超过并发上限直接 429，别让请求在内存里堆成雪崩 let _permit = w.inflight.clone().try_acquire_owned() .map_err(|_| (StatusCode::TOO_MANY_REQUESTS, "服务繁忙，请稍后重试".to_string()))?; // 2) 分词：这一步是纯 CPU，且 tokenizer 是 Sync 的，可以并发做 let p = prepare(&w.tokenizer, &req.text) .map_err(|e| (StatusCode::BAD_REQUEST, e.to_string()))?; // 3) 投递到推理队列，等结果 let (tx, rx) = oneshot::channel(); w.tx.send(Job { ids: p.ids, mask: p.mask, reply: tx }).await .map_err(|_| (StatusCode::SERVICE_UNAVAILABLE, "推理 worker 不可用".to_string()))?; // 4) 给推理加超时：模型卡住时不能让请求无限等 let (class_id, score) = tokio::time::timeout(std::time::Duration::from_secs(10), rx) .await .map_err(|_| (StatusCode::GATEWAY_TIMEOUT, "推理超时".to_string()))? .map_err(|_| (StatusCode::INTERNAL_SERVER_ERROR, "worker 异常退出".to_string()))? .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?; Ok(Json(PredictResp { label: LABELS[class_id.min(LABELS.len() - 1)].to_string(), score, tokens: p.raw_len, truncated: p.truncated, latency_ms: started.elapsed().as_millis() as u64, })) } // 路由与状态装配 pub fn router(w: Arc<ModelWorker>) -> Router { Router::new() .route("/predict", post(predict)) .route("/healthz", axum::routing::get(|| async { "ok" })) .with_state(w) }

pub async fn run_worker(mut session: ort::session::Session, mut rx: mpsc::Receiver<Job>) { while let Some(job) = rx.recv().await { // 把两个长度不同的输入补齐到同一长度：单条推理也要保证形状一致 let len = job.ids.len().max(1); let ids = ndarray::Array2::from_shape_vec((1, len), { let mut v = job.ids.clone(); v.resize(len, 0); // pad_id = 0 v }).unwrap(); let mask = ndarray::Array2::from_shape_vec((1, len), { let mut v = job.mask.clone(); v.resize(len, 0); // padding 位置的 mask 必须是 0 v }).unwrap(); let result = session .run(ort::inputs![ "input_ids" => ort::value::TensorRef::from_array_view(&ids).unwrap(), "attention_mask" => ort::value::TensorRef::from_array_view(&mask).unwrap(), ]) .map_err(|e| anyhow::anyhow!("推理失败: {e}")) .and_then(|outputs| { let logits = outputs["logits"].try_extract_array::<f32>()?; let row: Vec<f32> = logits.index_axis(ndarray::Axis(0), 0).iter().copied().collect(); let probs = softmax(&row); let best = argmax(&probs); Ok((best, probs[best])) }); // 无论成功失败都要回信，否则调用方会一直等到超时 let _ = job.reply.send(result); } // 通道关闭（发送方全部 drop）时退出循环，Session 在这里被释放 tracing::info!("推理 worker 退出"); }

use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::Arc; use tokio::sync::{Mutex, RwLock}; /// 一个可热插拔的模型持有者 pub struct ModelSlot { // 读多写少：用 RwLock，请求路径只需要 read + 更新一下时间戳 inner: RwLock<Option<LoadedModel>>, // 上次使用时刻（Unix 秒）。用原子变量避免请求路径写锁 last_used: AtomicU64, // 加载锁：保证"并发首次请求"只有一个真的去加载 loading: Mutex<()>, model_path: String, idle_ttl_secs: u64, } pub struct LoadedModel { pub session: ort::session::Session, pub tokenizer: Arc<Tokenizer>, } impl ModelSlot { pub fn new(model_path: String, idle_ttl_secs: u64) -> Self { Self { inner: RwLock::new(None), // 初始设为 0，保证第一轮扫描就会把它当作"空闲"（此时本来也没加载） last_used: AtomicU64::new(0), loading: Mutex::new(()), model_path, idle_ttl_secs, } } /// 取模型：已加载就直接用；没加载就（在锁保护下）加载一次 pub async fn acquire(&self) -> anyhow::Result<()> { // 快路径：读锁检查，命中就只更新一下时间戳 { let guard = self.inner.read().await; if guard.is_some() { self.touch(); return Ok(()); } } // 慢路径：加"加载锁"串行化。多个并发请求只会有第一个真的加载， // 其余的拿到锁之后会重新检查（double-check），发现已加载就直接返回 let _g = self.loading.lock().await; { let guard = self.inner.read().await; if guard.is_some() { self.touch(); return Ok(()); } } // 真正的加载。注意这里没有持有任何锁（除了 loading），不会阻塞其他模型 tracing::info!(path = %self.model_path, "开始加载模型（冷启动）"); let started = std::time::Instant::now(); let tokenizer = Arc::new(load_tokenizer(std::path::Path::new(&self.model_path))?); let session = build_session(&format!("{}/model.onnx", self.model_path), 4)?; tracing::info!(elapsed_ms = started.elapsed().as_millis() as u64, "模型加载完成"); *self.inner.write().await = Some(LoadedModel { session, tokenizer }); self.touch(); Ok(()) } fn touch(&self) { let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); // Relaxed 就够了：这里只是一个"大致时间"，不需要跨线程的严格顺序 self.last_used.store(now, Ordering::Relaxed); } /// 空闲判定 + 卸载 pub async fn maybe_unload(&self) -> bool { let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH).unwrap().as_secs(); let last = self.last_used.load(Ordering::Relaxed); // last == 0 表示从未使用过，此时本来就没加载，不用处理 if last == 0 || now.saturating_sub(last) < self.idle_ttl_secs { return false; } let mut guard = self.inner.write().await; if guard.is_some() { // 直接置 None，Session 与 Tokenizer 会在这里被 Drop，内存立刻归还 *guard = None; // 关键：卸载后把时间戳归零，避免下一轮扫描重复触发 self.last_used.store(0, Ordering::Relaxed); tracing::info!(path = %self.model_path, "空闲超时，模型已卸载"); return true; } false } } /// 后台扫描任务：每 10 秒检查一次，超过 idle_ttl_secs 未使用就卸载 pub fn spawn_idle_reaper(slots: Vec<Arc<ModelSlot>>) -> tokio::task::JoinHandle<()> { tokio::spawn(async move { // 扫描间隔要比 TTL 小得多，否则实际空闲时间会长出"一个扫描周期" let mut ticker = tokio::time::interval(std::time::Duration::from_secs(10)); loop { ticker.tick().await; for slot in &slots { if slot.maybe_unload().await { // 卸载动作做成指标，能看出"是不是在被反复加载卸载" tracing::info!("idle_unload_total", "模型空闲卸载计数 +1"); } } } }) }