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`, `mistralrs.rs`, 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 ocr;
35mod presets;
36mod raw;
37mod rerank;
38mod rerank_generative;
39
40use std::path::PathBuf;
41use std::sync::Arc;
42
43use async_trait::async_trait;
44use dashmap::DashMap;
45
46use crate::api::{ModelAliasSpec, ModelTask};
47use crate::error::{Result, RuntimeError};
48use crate::traits::{
49    DocumentExtractionModel, EmbeddingModel, HybridEmbeddingModel, ImageEmbeddingModel,
50    LoadedModelHandle, ModelProvider, MultiVectorEmbeddingModel, NlpModel, OcrModel,
51    ProviderCapabilities, ProviderHealth, RawTensorModel, RerankerModel, SparseEmbeddingModel,
52};
53
54pub struct LocalOnnxProvider {
55    base_dir: Option<PathBuf>,
56    sessions: DashMap<String, Arc<raw::LoadedOnnxSession>>,
57}
58
59impl Default for LocalOnnxProvider {
60    fn default() -> Self {
61        Self::new()
62    }
63}
64
65impl LocalOnnxProvider {
66    pub fn new() -> Self {
67        Self {
68            base_dir: None,
69            sessions: DashMap::new(),
70        }
71    }
72
73    pub fn with_base_dir(mut self, base_dir: impl Into<PathBuf>) -> Self {
74        self.base_dir = Some(base_dir.into());
75        self
76    }
77}
78
79#[async_trait]
80impl ModelProvider for LocalOnnxProvider {
81    fn provider_id(&self) -> &'static str {
82        "local/onnx"
83    }
84
85    fn capabilities(&self) -> ProviderCapabilities {
86        ProviderCapabilities {
87            supported_tasks: vec![
88                ModelTask::Raw,
89                ModelTask::Rerank,
90                ModelTask::Embed,
91                ModelTask::Nlp,
92                ModelTask::EmbedImage,
93                ModelTask::Ocr,
94                ModelTask::DocumentExtract,
95                ModelTask::EmbedSparse,
96                ModelTask::EmbedMultiVector,
97                ModelTask::EmbedHybrid,
98            ],
99        }
100    }
101
102    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
103        match spec.task {
104            ModelTask::Raw => {
105                let runner: Arc<dyn RawTensorModel> =
106                    raw::load_raw(spec, self.base_dir.as_deref(), &self.sessions).await?;
107                Ok(Arc::new(runner) as LoadedModelHandle)
108            }
109            ModelTask::Rerank => {
110                let style = spec
111                    .options
112                    .get("style")
113                    .and_then(|v| v.as_str())
114                    .unwrap_or("cross-encoder");
115                let reranker: Arc<dyn RerankerModel> = match style {
116                    "cross-encoder" => rerank::load_rerank(spec).await?,
117                    "generative" => rerank_generative::load(spec).await?,
118                    other => {
119                        return Err(RuntimeError::Config(format!(
120                            "Unknown rerank style '{other}' for alias '{}'; expected one of cross-encoder, generative",
121                            spec.alias
122                        )));
123                    }
124                };
125                Ok(Arc::new(reranker) as LoadedModelHandle)
126            }
127            ModelTask::Embed => {
128                let embedder: Arc<dyn EmbeddingModel> = embedding::load_embedding(spec).await?;
129                Ok(Arc::new(embedder) as LoadedModelHandle)
130            }
131            ModelTask::Nlp => {
132                let model: Arc<dyn NlpModel> = nlp::load_nlp(spec).await?;
133                Ok(Arc::new(model) as LoadedModelHandle)
134            }
135            ModelTask::EmbedImage => {
136                let model: Arc<dyn ImageEmbeddingModel> =
137                    image_embed::load_image_embedder(spec).await?;
138                Ok(Arc::new(model) as LoadedModelHandle)
139            }
140            ModelTask::Ocr => {
141                let model: Arc<dyn OcrModel> = ocr::load_ocr(spec).await?;
142                Ok(Arc::new(model) as LoadedModelHandle)
143            }
144            ModelTask::DocumentExtract => {
145                let model: Arc<dyn DocumentExtractionModel> =
146                    document_extract::load_document_extractor(spec).await?;
147                Ok(Arc::new(model) as LoadedModelHandle)
148            }
149            ModelTask::EmbedSparse => {
150                let model: Arc<dyn SparseEmbeddingModel> =
151                    embed_sparse::load_sparse_embedder(spec).await?;
152                Ok(Arc::new(model) as LoadedModelHandle)
153            }
154            ModelTask::EmbedMultiVector => {
155                let model: Arc<dyn MultiVectorEmbeddingModel> =
156                    embed_multi_vector::load_multi_vector_embedder(spec).await?;
157                Ok(Arc::new(model) as LoadedModelHandle)
158            }
159            ModelTask::EmbedHybrid => {
160                let model: Arc<dyn HybridEmbeddingModel> =
161                    embed_hybrid::load_hybrid_embedder(spec).await?;
162                Ok(Arc::new(model) as LoadedModelHandle)
163            }
164            _ => Err(RuntimeError::CapabilityMismatch(format!(
165                "ONNX provider does not support task {:?}",
166                spec.task
167            ))),
168        }
169    }
170
171    async fn health(&self) -> ProviderHealth {
172        ProviderHealth::Healthy
173    }
174}
uni_xervo/provider/local_onnx.rs

uni_xervo/provider/
local_onnx.rs
