gte/params/
mod.rs

1//! Parameters for embedding or re-ranking
2
3use crate::embed::output::{ExtractorMode, OutputId};
4
5/// Parameters
6/// 
7/// Prefer using `default()` and set individual parameters as needed.
8pub struct Parameters {
9    max_length: Option<usize>,
10    sigmoid: bool,
11    output_id: OutputId,
12    mode: ExtractorMode,
13}
14
15impl Parameters {
16    
17    /// Input truncation (nb of tokens)
18    pub fn with_max_length(mut self, max_length: Option<usize>) -> Self {
19        self.max_length = max_length;
20        self
21    }
22
23    /// Input truncation (nb of tokens)
24    pub fn max_length(&self) -> Option<usize> {
25        self.max_length
26    }
27
28    /// Apply sigmoid (for re-reanking)
29    pub fn with_sigmoid(mut self, sigmoid: bool) -> Self {
30        self.sigmoid = sigmoid;
31        self
32    }
33
34    /// Apply sigmoid (for re-reanking)
35    pub fn sigmoid(&self) -> bool {
36        self.sigmoid
37    }
38
39    /// Set output tensor identifier (eg. `last_hidden_state`)
40    pub fn with_output_id(mut self, id: &str) -> Self {
41        self.output_id = id.into();
42        self
43    }
44
45    /// Output tensor identifier
46    pub fn output_id(&self) -> &OutputId {
47        &self.output_id
48    }
49
50    /// Set embeddings extraction mode
51    pub fn with_mode(mut self, mode: ExtractorMode) -> Self {
52        self.mode = mode;
53        self
54    }
55
56    /// Embeddings extraction mode
57    pub fn mode(&self) -> ExtractorMode {
58        self.mode
59    }
60    
61}
62
63
64impl Default for Parameters {
65    fn default() -> Self {
66        Self { 
67            max_length: None,
68            sigmoid: false,
69            output_id: OutputId::default(),
70            mode: ExtractorMode::default(),
71        }
72    }
73}
gte/params/mod.rs

gte/params/
mod.rs
