uni_xervo/provider/
local_onnx.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright 2024-2026 Dragonscale Team
3
4//! Local ONNX Runtime provider (`local/onnx`).
5//!
6//! Single backend, multiple tasks. The provider declares
7//! [`ModelTask::Raw`] and [`ModelTask::Rerank`] in [`ProviderCapabilities`]
8//! and dispatches inside [`ModelProvider::load`] — matching the pattern
9//! used by `cohere.rs`, `llamacpp/`, etc.
10//!
11//! Task implementations live in private submodules:
12//!
13//! - `raw` — arbitrary ONNX tensor execution via [`RawTensorModel`].
14//! - `rerank` — cross-encoder reranking via [`RerankerModel`].
15//! - `embedding` — dense text embeddings via [`EmbeddingModel`].
16
17// Reserved for downstream provider impls (Granite-Docling / MinerU /
18// olmOCR / future decoder LMs). The module is built and unit-tested
19// independently; #[allow(dead_code)] suppresses the unused-warning
20// noise until the first consumer wires it.
21#[allow(dead_code)]
22mod autoreg;
23mod common;
24mod decoder_inputs;
25mod det;
26mod document_extract;
27mod embed_hybrid;
28mod embed_multi_vector;
29mod embed_sparse;
30mod embedding;
31mod image;
32mod image_embed;
33mod nlp;
34mod nlp_v7;
35mod ocr;
36mod presets;
37mod raw;
38mod rerank;
39mod rerank_generative;
40
41use std::path::PathBuf;
42use std::sync::Arc;
43
44use async_trait::async_trait;
45use dashmap::DashMap;
46
47use crate::api::{ModelAliasSpec, ModelTask};
48use crate::error::{Result, RuntimeError};
49use crate::traits::{
50    DocumentExtractionModel, EmbeddingModel, HybridEmbeddingModel, ImageEmbeddingModel,
51    LoadedModelHandle, ModelProvider, MultiVectorEmbeddingModel, NlpModel, OcrModel,
52    ProviderCapabilities, ProviderHealth, RawTensorModel, RerankerModel, SparseEmbeddingModel,
53};
54
55pub struct LocalOnnxProvider {
56    base_dir: Option<PathBuf>,
57    sessions: DashMap<String, Arc<raw::LoadedOnnxSession>>,
58}
59
60impl Default for LocalOnnxProvider {
61    fn default() -> Self {
62        Self::new()
63    }
64}
65
66impl LocalOnnxProvider {
67    pub fn new() -> Self {
68        Self {
69            base_dir: None,
70            sessions: DashMap::new(),
71        }
72    }
73
74    pub fn with_base_dir(mut self, base_dir: impl Into<PathBuf>) -> Self {
75        self.base_dir = Some(base_dir.into());
76        self
77    }
78}
79
80#[async_trait]
81impl ModelProvider for LocalOnnxProvider {
82    fn provider_id(&self) -> &'static str {
83        "local/onnx"
84    }
85
86    fn capabilities(&self) -> ProviderCapabilities {
87        ProviderCapabilities {
88            supported_tasks: vec![
89                ModelTask::Raw,
90                ModelTask::Rerank,
91                ModelTask::Embed,
92                ModelTask::Nlp,
93                ModelTask::EmbedImage,
94                ModelTask::Ocr,
95                ModelTask::DocumentExtract,
96                ModelTask::EmbedSparse,
97                ModelTask::EmbedMultiVector,
98                ModelTask::EmbedHybrid,
99            ],
100        }
101    }
102
103    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
104        match spec.task {
105            ModelTask::Raw => {
106                let runner: Arc<dyn RawTensorModel> =
107                    raw::load_raw(spec, self.base_dir.as_deref(), &self.sessions).await?;
108                Ok(Arc::new(runner) as LoadedModelHandle)
109            }
110            ModelTask::Rerank => {
111                let style = spec
112                    .options
113                    .get("style")
114                    .and_then(|v| v.as_str())
115                    .unwrap_or("cross-encoder");
116                let reranker: Arc<dyn RerankerModel> = match style {
117                    "cross-encoder" => rerank::load_rerank(spec).await?,
118                    "generative" => rerank_generative::load(spec).await?,
119                    other => {
120                        return Err(RuntimeError::Config(format!(
121                            "Unknown rerank style '{other}' for alias '{}'; expected one of cross-encoder, generative",
122                            spec.alias
123                        )));
124                    }
125                };
126                Ok(Arc::new(reranker) as LoadedModelHandle)
127            }
128            ModelTask::Embed => {
129                let embedder: Arc<dyn EmbeddingModel> = embedding::load_embedding(spec).await?;
130                Ok(Arc::new(embedder) as LoadedModelHandle)
131            }
132            ModelTask::Nlp => {
133                // Two generations of the kniv NLP family share this task.
134                // The v1 cascade is one graph of 128 subwords; v7 is eight
135                // graphs over word-level windows, with coreference, typed
136                // entities and schema relations. They have no code in
137                // common beyond this dispatch.
138                let model: Arc<dyn NlpModel> = match nlp_v7::detect_generation(spec).await? {
139                    nlp_v7::Generation::V7 => nlp_v7::load_nlp_v7(spec).await?,
140                    nlp_v7::Generation::V1 => nlp::load_nlp(spec).await?,
141                };
142                Ok(Arc::new(model) as LoadedModelHandle)
143            }
144            ModelTask::EmbedImage => {
145                let model: Arc<dyn ImageEmbeddingModel> =
146                    image_embed::load_image_embedder(spec).await?;
147                Ok(Arc::new(model) as LoadedModelHandle)
148            }
149            ModelTask::Ocr => {
150                let model: Arc<dyn OcrModel> = ocr::load_ocr(spec).await?;
151                Ok(Arc::new(model) as LoadedModelHandle)
152            }
153            ModelTask::DocumentExtract => {
154                let model: Arc<dyn DocumentExtractionModel> =
155                    document_extract::load_document_extractor(spec).await?;
156                Ok(Arc::new(model) as LoadedModelHandle)
157            }
158            ModelTask::EmbedSparse => {
159                let model: Arc<dyn SparseEmbeddingModel> =
160                    embed_sparse::load_sparse_embedder(spec).await?;
161                Ok(Arc::new(model) as LoadedModelHandle)
162            }
163            ModelTask::EmbedMultiVector => {
164                let model: Arc<dyn MultiVectorEmbeddingModel> =
165                    embed_multi_vector::load_multi_vector_embedder(spec).await?;
166                Ok(Arc::new(model) as LoadedModelHandle)
167            }
168            ModelTask::EmbedHybrid => {
169                let model: Arc<dyn HybridEmbeddingModel> =
170                    embed_hybrid::load_hybrid_embedder(spec).await?;
171                Ok(Arc::new(model) as LoadedModelHandle)
172            }
173            _ => Err(RuntimeError::CapabilityMismatch(format!(
174                "ONNX provider does not support task {:?}",
175                spec.task
176            ))),
177        }
178    }
179
180    async fn health(&self) -> ProviderHealth {
181        ProviderHealth::Healthy
182    }
183}
uni_xervo/provider/local_onnx.rs

uni_xervo/provider/
local_onnx.rs
