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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
use crate::get_cache_dir;
use ort::execution_providers::ExecutionProviderDispatch;
use std::path::PathBuf;
pub trait HasMaxLength {
const MAX_LENGTH: usize;
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct InitOptionsWithLength<M> {
pub model_name: M,
pub execution_providers: Vec<ExecutionProviderDispatch>,
pub cache_dir: PathBuf,
pub show_download_progress: bool,
pub max_length: usize,
/// Number of intra-op threads for ONNX Runtime. `None` (the default) uses
/// every available CPU core via `std::thread::available_parallelism`.
/// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
pub intra_threads: Option<usize>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}
#[derive(Debug, Clone)]
#[non_exhaustive]
pub struct InitOptions<M> {
pub model_name: M,
pub execution_providers: Vec<ExecutionProviderDispatch>,
pub cache_dir: PathBuf,
pub show_download_progress: bool,
/// Number of intra-op threads for ONNX Runtime. `None` (the default) uses
/// every available CPU core via `std::thread::available_parallelism`.
/// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
pub intra_threads: Option<usize>,
/// ONNX Runtime session configuration entries, applied with
/// `SessionBuilder::with_config_entry`. Use this for settings that have
/// no dedicated builder method, such as `mlas.disable_kleidiai`.
pub session_config: Vec<(String, String)>,
}
impl<M: Default + HasMaxLength> Default for InitOptionsWithLength<M> {
fn default() -> Self {
Self {
model_name: M::default(),
execution_providers: Default::default(),
cache_dir: get_cache_dir().into(),
show_download_progress: true,
max_length: M::MAX_LENGTH,
intra_threads: None,
session_config: Vec::new(),
}
}
}
impl<M: Default> Default for InitOptions<M> {
fn default() -> Self {
Self {
model_name: M::default(),
execution_providers: Default::default(),
cache_dir: get_cache_dir().into(),
show_download_progress: true,
intra_threads: None,
session_config: Vec::new(),
}
}
}
impl<M: Default + HasMaxLength> InitOptionsWithLength<M> {
/// Create a new InitOptionsWithLength with the given model name
pub fn new(model_name: M) -> Self {
Self {
model_name,
..Default::default()
}
}
/// Set the maximum length
pub fn with_max_length(mut self, max_length: usize) -> Self {
self.max_length = max_length;
self
}
/// Set the cache directory for the model file
pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self {
self.cache_dir = cache_dir;
self
}
/// Set the execution providers for the model
pub fn with_execution_providers(
mut self,
execution_providers: Vec<ExecutionProviderDispatch>,
) -> Self {
self.execution_providers = execution_providers;
self
}
/// Set the number of intra-op threads ONNX Runtime uses. By default
/// (`None`) all available CPU cores are used; capping this limits CPU
/// usage at the cost of per-inference throughput.
pub fn with_intra_threads(mut self, intra_threads: usize) -> Self {
self.intra_threads = Some(intra_threads);
self
}
/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}
/// Set whether to show download progress
pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
self.show_download_progress = show_download_progress;
self
}
}
impl<M: Default> InitOptions<M> {
/// Create a new InitOptions with the given model name
pub fn new(model_name: M) -> Self {
Self {
model_name,
..Default::default()
}
}
/// Set the cache directory for the model file
pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self {
self.cache_dir = cache_dir;
self
}
/// Set the execution providers for the model
pub fn with_execution_providers(
mut self,
execution_providers: Vec<ExecutionProviderDispatch>,
) -> Self {
self.execution_providers = execution_providers;
self
}
/// Set the number of intra-op threads ONNX Runtime uses. By default
/// (`None`) all available CPU cores are used; capping this limits CPU
/// usage at the cost of per-inference throughput.
pub fn with_intra_threads(mut self, intra_threads: usize) -> Self {
self.intra_threads = Some(intra_threads);
self
}
/// Add an ONNX Runtime session configuration entry, applied with
/// `SessionBuilder::with_config_entry`. Call it once per entry.
/// Example: `.with_session_config("mlas.disable_kleidiai", "1")`.
pub fn with_session_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.session_config.push((key.into(), value.into()));
self
}
/// Set whether to show download progress
pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
self.show_download_progress = show_download_progress;
self
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn intra_threads_defaults_none_and_builder_sets() {
let o = InitOptions::<crate::ImageEmbeddingModel>::default();
assert_eq!(o.intra_threads, None);
let o = o.with_intra_threads(4);
assert_eq!(o.intra_threads, Some(4));
}
}