fastembed/reranking/
init.rs

1use super::DEFAULT_MAX_LENGTH;
2use crate::{
3    init::{HasMaxLength, InitOptionsWithLength},
4    RerankerModel, TokenizerFiles,
5};
6use ort::{execution_providers::ExecutionProviderDispatch, session::Session};
7use std::path::PathBuf;
8use tokenizers::Tokenizer;
9
10#[derive(Debug)]
11pub struct TextRerank {
12    pub tokenizer: Tokenizer,
13    pub(crate) session: Session,
14    pub(crate) need_token_type_ids: bool,
15}
16
17impl HasMaxLength for RerankerModel {
18    const MAX_LENGTH: usize = DEFAULT_MAX_LENGTH;
19}
20
21/// Options for initializing the reranking models
22pub type RerankInitOptions = InitOptionsWithLength<RerankerModel>;
23
24/// Options for initializing UserDefinedRerankerModel
25///
26/// Model files are held by the UserDefinedRerankerModel struct
27#[derive(Debug, Clone)]
28#[non_exhaustive]
29pub struct RerankInitOptionsUserDefined {
30    pub execution_providers: Vec<ExecutionProviderDispatch>,
31    pub max_length: usize,
32}
33
34impl Default for RerankInitOptionsUserDefined {
35    fn default() -> Self {
36        Self {
37            execution_providers: Default::default(),
38            max_length: DEFAULT_MAX_LENGTH,
39        }
40    }
41}
42
43/// Convert RerankInitOptions to RerankInitOptionsUserDefined
44///
45/// This is useful for when the user wants to use the same options for both the default and user-defined models
46impl From<RerankInitOptions> for RerankInitOptionsUserDefined {
47    fn from(options: RerankInitOptions) -> Self {
48        RerankInitOptionsUserDefined {
49            execution_providers: options.execution_providers,
50            max_length: options.max_length,
51        }
52    }
53}
54
55/// Enum for the source of the onnx file
56///
57/// User-defined models can either be in memory or on disk
58#[derive(Debug, Clone, PartialEq, Eq)]
59pub enum OnnxSource {
60    Memory(Vec<u8>),
61    File(PathBuf),
62}
63
64impl From<Vec<u8>> for OnnxSource {
65    fn from(bytes: Vec<u8>) -> Self {
66        OnnxSource::Memory(bytes)
67    }
68}
69
70impl From<PathBuf> for OnnxSource {
71    fn from(path: PathBuf) -> Self {
72        OnnxSource::File(path)
73    }
74}
75
76/// Struct for "bring your own" reranking models
77///
78/// The onnx_file and tokenizer_files are expecting the files' bytes
79#[derive(Debug, Clone, PartialEq, Eq)]
80#[non_exhaustive]
81pub struct UserDefinedRerankingModel {
82    pub onnx_source: OnnxSource,
83    pub tokenizer_files: TokenizerFiles,
84}
85
86impl UserDefinedRerankingModel {
87    pub fn new(onnx_source: impl Into<OnnxSource>, tokenizer_files: TokenizerFiles) -> Self {
88        Self {
89            onnx_source: onnx_source.into(),
90            tokenizer_files,
91        }
92    }
93}
94
95/// Rerank result.
96#[derive(Debug, PartialEq, Clone)]
97pub struct RerankResult {
98    pub document: Option<String>,
99    pub score: f32,
100    pub index: usize,
101}
fastembed/reranking/init.rs

fastembed/reranking/
init.rs
