fastembed/
init.rs

1use crate::get_cache_dir;
2use ort::execution_providers::ExecutionProviderDispatch;
3use std::path::PathBuf;
4
5pub trait HasMaxLength {
6    const MAX_LENGTH: usize;
7}
8
9#[derive(Debug, Clone)]
10#[non_exhaustive]
11pub struct InitOptionsWithLength<M> {
12    pub model_name: M,
13    pub execution_providers: Vec<ExecutionProviderDispatch>,
14    pub cache_dir: PathBuf,
15    pub show_download_progress: bool,
16    pub max_length: usize,
17    /// Number of intra-op threads for ONNX Runtime. `None` (the default) uses
18    /// every available CPU core via `std::thread::available_parallelism`.
19    /// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
20    pub intra_threads: Option<usize>,
21}
22
23#[derive(Debug, Clone)]
24#[non_exhaustive]
25pub struct InitOptions<M> {
26    pub model_name: M,
27    pub execution_providers: Vec<ExecutionProviderDispatch>,
28    pub cache_dir: PathBuf,
29    pub show_download_progress: bool,
30    /// Number of intra-op threads for ONNX Runtime. `None` (the default) uses
31    /// every available CPU core via `std::thread::available_parallelism`.
32    /// Set this to cap CPU usage (e.g. on laptops) at the cost of throughput.
33    pub intra_threads: Option<usize>,
34}
35
36impl<M: Default + HasMaxLength> Default for InitOptionsWithLength<M> {
37    fn default() -> Self {
38        Self {
39            model_name: M::default(),
40            execution_providers: Default::default(),
41            cache_dir: get_cache_dir().into(),
42            show_download_progress: true,
43            max_length: M::MAX_LENGTH,
44            intra_threads: None,
45        }
46    }
47}
48
49impl<M: Default> Default for InitOptions<M> {
50    fn default() -> Self {
51        Self {
52            model_name: M::default(),
53            execution_providers: Default::default(),
54            cache_dir: get_cache_dir().into(),
55            show_download_progress: true,
56            intra_threads: None,
57        }
58    }
59}
60
61impl<M: Default + HasMaxLength> InitOptionsWithLength<M> {
62    /// Create a new InitOptionsWithLength with the given model name
63    pub fn new(model_name: M) -> Self {
64        Self {
65            model_name,
66            ..Default::default()
67        }
68    }
69
70    /// Set the maximum length
71    pub fn with_max_length(mut self, max_length: usize) -> Self {
72        self.max_length = max_length;
73        self
74    }
75
76    /// Set the cache directory for the model file
77    pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self {
78        self.cache_dir = cache_dir;
79        self
80    }
81
82    /// Set the execution providers for the model
83    pub fn with_execution_providers(
84        mut self,
85        execution_providers: Vec<ExecutionProviderDispatch>,
86    ) -> Self {
87        self.execution_providers = execution_providers;
88        self
89    }
90
91    /// Set the number of intra-op threads ONNX Runtime uses. By default
92    /// (`None`) all available CPU cores are used; capping this limits CPU
93    /// usage at the cost of per-inference throughput.
94    pub fn with_intra_threads(mut self, intra_threads: usize) -> Self {
95        self.intra_threads = Some(intra_threads);
96        self
97    }
98
99    /// Set whether to show download progress
100    pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
101        self.show_download_progress = show_download_progress;
102        self
103    }
104}
105
106impl<M: Default> InitOptions<M> {
107    /// Create a new InitOptions with the given model name
108    pub fn new(model_name: M) -> Self {
109        Self {
110            model_name,
111            ..Default::default()
112        }
113    }
114
115    /// Set the cache directory for the model file
116    pub fn with_cache_dir(mut self, cache_dir: PathBuf) -> Self {
117        self.cache_dir = cache_dir;
118        self
119    }
120
121    /// Set the execution providers for the model
122    pub fn with_execution_providers(
123        mut self,
124        execution_providers: Vec<ExecutionProviderDispatch>,
125    ) -> Self {
126        self.execution_providers = execution_providers;
127        self
128    }
129
130    /// Set the number of intra-op threads ONNX Runtime uses. By default
131    /// (`None`) all available CPU cores are used; capping this limits CPU
132    /// usage at the cost of per-inference throughput.
133    pub fn with_intra_threads(mut self, intra_threads: usize) -> Self {
134        self.intra_threads = Some(intra_threads);
135        self
136    }
137
138    /// Set whether to show download progress
139    pub fn with_show_download_progress(mut self, show_download_progress: bool) -> Self {
140        self.show_download_progress = show_download_progress;
141        self
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn intra_threads_defaults_none_and_builder_sets() {
151        let o = InitOptions::<crate::ImageEmbeddingModel>::default();
152        assert_eq!(o.intra_threads, None);
153        let o = o.with_intra_threads(4);
154        assert_eq!(o.intra_threads, Some(4));
155    }
156}
fastembed/init.rs

fastembed/
init.rs
