ferrum_engine/
embedding_engine.rs

1//! Lightweight engine for embedding models (CLIP, BERT, etc.).
2//!
3//! Wraps a `ClipModelExecutor` and implements `EmbedEngine` directly —
4//! no LLM-method stubs, no fake `InferenceEngine::infer` impls. The
5//! `InferenceEngine` supertrait covers lifecycle/status; the modality
6//! itself goes through `EmbedEngine`.
7
8use async_trait::async_trait;
9use ferrum_interfaces::engine::{EmbedEngine, InferenceEngine};
10use ferrum_models::ClipModelExecutor;
11use ferrum_types::{EngineConfig, EngineMetrics, EngineStatus, FerrumError, Result};
12use std::sync::Arc;
13
14/// Embedding-only engine wrapping a ClipModelExecutor.
15pub struct EmbeddingEngine {
16    executor: Arc<ClipModelExecutor>,
17    tokenizer: Option<tokenizers::Tokenizer>,
18    config: EngineConfig,
19}
20
21impl EmbeddingEngine {
22    pub fn new(executor: ClipModelExecutor, config: EngineConfig) -> Self {
23        Self {
24            executor: Arc::new(executor),
25            tokenizer: None,
26            config,
27        }
28    }
29
30    /// Set tokenizer for text embedding.
31    pub fn with_tokenizer(mut self, tokenizer: tokenizers::Tokenizer) -> Self {
32        self.tokenizer = Some(tokenizer);
33        self
34    }
35}
36
37#[async_trait]
38impl InferenceEngine for EmbeddingEngine {
39    async fn status(&self) -> EngineStatus {
40        crate::modality_stubs::inert_status()
41    }
42
43    async fn shutdown(&self) -> Result<()> {
44        Ok(())
45    }
46
47    fn config(&self) -> &EngineConfig {
48        &self.config
49    }
50
51    fn metrics(&self) -> EngineMetrics {
52        crate::modality_stubs::inert_metrics()
53    }
54
55    async fn health_check(&self) -> ferrum_types::HealthStatus {
56        crate::modality_stubs::inert_health()
57    }
58}
59
60#[async_trait]
61impl EmbedEngine for EmbeddingEngine {
62    async fn embed_text(&self, text: &str) -> Result<Vec<f32>> {
63        let tokenizer = self
64            .tokenizer
65            .as_ref()
66            .ok_or_else(|| FerrumError::model("No tokenizer loaded for text embedding"))?;
67        let encoding = tokenizer
68            .encode(text, true)
69            .map_err(|e| FerrumError::model(format!("tokenize: {e}")))?;
70        let embedding = self.executor.embed_text(encoding.get_ids())?;
71        embedding
72            .squeeze(0)
73            .and_then(|t| t.to_dtype(candle_core::DType::F32))
74            .and_then(|t| t.to_vec1())
75            .map_err(|e| FerrumError::model(format!("embed_text tensor: {e}")))
76    }
77
78    async fn embed_image(&self, image: &str) -> Result<Vec<f32>> {
79        let embedding = if image.starts_with("data:") || image.len() > 1000 {
80            self.executor.embed_image_base64(image)?
81        } else {
82            self.executor.embed_image_path(image)?
83        };
84        embedding
85            .squeeze(0)
86            .and_then(|t| t.to_dtype(candle_core::DType::F32))
87            .and_then(|t| t.to_vec1())
88            .map_err(|e| FerrumError::model(format!("embed_image tensor: {e}")))
89    }
90
91    fn embedding_dim(&self) -> usize {
92        self.executor.projection_dim()
93    }
94}
ferrum_engine/embedding_engine.rs

ferrum_engine/
embedding_engine.rs
