aurora_semantic/embeddings/

providers.rs

//! Embedding providers - ONNX-based local model inference using ONNX Runtime.
//!
//! This module provides embedding generation using Microsoft's ONNX Runtime via the `ort` crate.
//! Users download ONNX models and tokenizers, then point to their location.
//!
//! # Supported Models
//!
//! Any ONNX model that outputs embeddings can be used:
//! - `jina-embeddings-v2-base-code` (recommended for code)
//! - `all-MiniLM-L6-v2`
//! - `bge-small-en-v1.5`
//! - Any sentence-transformer model exported to ONNX
//!
//! # Model Directory Structure
//!
//! ```text
//! model_dir/
//! ├── model.onnx          # ONNX model file (or model-w-mean-pooling.onnx)
//! ├── tokenizer.json      # HuggingFace tokenizer
//! └── config.json         # Optional: model config
//! ```
//!
//! # GPU Acceleration
//!
//! GPU support is available via Cargo features:
//! - `cuda` - NVIDIA CUDA (requires CUDA toolkit)
//! - `tensorrt` - NVIDIA TensorRT (requires TensorRT)
//! - `directml` - DirectML (Windows, AMD/Intel/NVIDIA)
//! - `coreml` - CoreML (macOS, Apple Silicon)
//!
//! Build with GPU support:
//! ```bash
//! cargo build --release --features cuda
//! ```

use std::path::Path;
use std::sync::Arc;

use ndarray::IxDyn;
use ort::session::Session;
use ort::value::Tensor;
use parking_lot::Mutex;
use serde::{Deserialize, Serialize};
use tokenizers::Tokenizer;

use crate::embeddings::Embedder;
use crate::error::{Error, Result};

/// Information about the execution provider (CPU/GPU) being used.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ExecutionProviderInfo {
    /// Name of the execution provider (e.g., "CPU", "CUDA", "DirectML", "CoreML").
    pub name: String,
    /// Whether GPU acceleration is being used.
    pub is_gpu: bool,
    /// Device ID if applicable.
    pub device_id: Option<u32>,
    /// Additional details about the provider.
    pub details: Option<String>,
}

impl ExecutionProviderInfo {
    /// Create CPU execution provider info.
    pub fn cpu() -> Self {
        Self {
            name: "CPU".to_string(),
            is_gpu: false,
            device_id: None,
            details: Some("Default CPU execution".to_string()),
        }
    }

    /// Create CUDA execution provider info.
    #[allow(dead_code)]
    pub fn cuda(device_id: u32) -> Self {
        Self {
            name: "CUDA".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("NVIDIA CUDA GPU (device {})", device_id)),
        }
    }

    /// Create TensorRT execution provider info.
    #[allow(dead_code)]
    pub fn tensorrt(device_id: u32) -> Self {
        Self {
            name: "TensorRT".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("NVIDIA TensorRT GPU (device {})", device_id)),
        }
    }

    /// Create DirectML execution provider info.
    #[allow(dead_code)]
    pub fn directml(device_id: u32) -> Self {
        Self {
            name: "DirectML".to_string(),
            is_gpu: true,
            device_id: Some(device_id),
            details: Some(format!("DirectML GPU (device {})", device_id)),
        }
    }

    /// Create CoreML execution provider info.
    #[allow(dead_code)]
    pub fn coreml() -> Self {
        Self {
            name: "CoreML".to_string(),
            is_gpu: true,
            device_id: None,
            details: Some("Apple CoreML (Neural Engine/GPU)".to_string()),
        }
    }

    /// Get a human-readable description.
    pub fn description(&self) -> String {
        if self.is_gpu {
            format!("{} (GPU accelerated)", self.name)
        } else {
            format!("{} (no GPU)", self.name)
        }
    }
}

/// ONNX-based embedding model using ONNX Runtime.
///
/// Loads an ONNX model and tokenizer from a directory and generates embeddings.
///
/// # Example
///
/// ```rust,ignore
/// use aurora_semantic::OnnxEmbedder;
///
/// // Point to your downloaded model directory
/// let embedder = OnnxEmbedder::from_directory("./models/jina-code")?;
///
/// let embedding = embedder.embed("fn main() { println!(\"Hello\"); }")?;
/// println!("Embedding dimension: {}", embedding.len());
/// ```
pub struct OnnxEmbedder {
    session: Arc<Mutex<Session>>,
    tokenizer: Tokenizer,
    dimension: usize,
    max_length: usize,
    execution_provider: ExecutionProviderInfo,
}

impl OnnxEmbedder {
    /// Load an ONNX model from a directory.
    ///
    /// The directory must contain:
    /// - An ONNX model file (model.onnx, model-w-mean-pooling.onnx, etc.)
    /// - `tokenizer.json` - HuggingFace tokenizer file
    pub fn from_directory<P: AsRef<Path>>(model_dir: P) -> Result<Self> {
        let model_dir = model_dir.as_ref();

        // Check for model file with various common names
        let model_names = [
            "model.onnx",
            "model_optimized.onnx",
            "model-w-mean-pooling.onnx",
            "model_quantized.onnx",
            "encoder_model.onnx",
        ];

        let model_path = model_names
            .iter()
            .map(|name| model_dir.join(name))
            .find(|p| p.exists())
            .ok_or_else(|| {
                Error::model_load(format!(
                    "No ONNX model file found in {}. Expected one of: {:?}",
                    model_dir.display(),
                    model_names
                ))
            })?;

        // Check tokenizer exists
        let tokenizer_path = model_dir.join("tokenizer.json");
        if !tokenizer_path.exists() {
            return Err(Error::model_load(format!(
                "tokenizer.json not found in {}",
                model_dir.display()
            )));
        }

        Self::new(&model_path, &tokenizer_path, None, 512)
    }

    /// Create a new ONNX embedder with explicit paths.
    pub fn new<P: AsRef<Path>>(
        model_path: P,
        tokenizer_path: P,
        dimension: Option<usize>,
        max_length: usize,
    ) -> Result<Self> {
        let model_path = model_path.as_ref();
        let tokenizer_path = tokenizer_path.as_ref();

        // Load tokenizer
        let tokenizer = Tokenizer::from_file(tokenizer_path)
            .map_err(|e| Error::model_load(format!("Failed to load tokenizer: {}", e)))?;

        // Load ONNX model with ort
        tracing::info!("Loading ONNX model from: {}", model_path.display());

        #[allow(unused_mut)]
        let mut builder = Session::builder()
            .map_err(|e| Error::model_load(format!("Failed to create session builder: {}", e)))?;

        // Track which execution provider we're using
        #[allow(unused_mut)]
        let mut execution_provider = ExecutionProviderInfo::cpu();

        // Configure execution providers based on available features
        #[cfg(feature = "cuda")]
        {
            use ort::execution_providers::CUDAExecutionProvider;
            tracing::info!("CUDA support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([CUDAExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure CUDA: {}", e)))?;
            execution_provider = ExecutionProviderInfo::cuda(0);
        }

        #[cfg(feature = "tensorrt")]
        {
            use ort::execution_providers::TensorRTExecutionProvider;
            tracing::info!("TensorRT support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([TensorRTExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure TensorRT: {}", e)))?;
            execution_provider = ExecutionProviderInfo::tensorrt(0);
        }

        #[cfg(feature = "directml")]
        {
            use ort::execution_providers::DirectMLExecutionProvider;
            tracing::info!("DirectML support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([DirectMLExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure DirectML: {}", e)))?;
            execution_provider = ExecutionProviderInfo::directml(0);
        }

        #[cfg(feature = "coreml")]
        {
            use ort::execution_providers::CoreMLExecutionProvider;
            tracing::info!("CoreML support enabled, attempting GPU acceleration");
            builder = builder
                .with_execution_providers([CoreMLExecutionProvider::default().build()])
                .map_err(|e| Error::model_load(format!("Failed to configure CoreML: {}", e)))?;
            execution_provider = ExecutionProviderInfo::coreml();
        }

        let session = builder
            .with_intra_threads(4)
            .map_err(|e| Error::model_load(format!("Failed to set threads: {}", e)))?
            .commit_from_file(model_path)
            .map_err(|e| Error::model_load(format!("Failed to load ONNX model: {}", e)))?;

        // Default dimension - will be detected from actual output
        let dimension = dimension.unwrap_or(768);

        tracing::info!(
            "Loaded ONNX model (dim={}, max_len={}, provider={})",
            dimension,
            max_length,
            execution_provider.description()
        );

        Ok(Self {
            session: Arc::new(Mutex::new(session)),
            tokenizer,
            dimension,
            max_length,
            execution_provider,
        })
    }

    /// Set the maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.max_length = max_length;
        self
    }

    /// Get information about the execution provider (CPU/GPU).
    pub fn execution_provider(&self) -> &ExecutionProviderInfo {
        &self.execution_provider
    }

    /// Check if GPU acceleration is being used.
    pub fn is_gpu_accelerated(&self) -> bool {
        self.execution_provider.is_gpu
    }

    /// Mean pooling over token embeddings.
    fn mean_pooling(&self, data: &[f32], shape: &[i64], attention_mask: &[i64]) -> Vec<f32> {
        if shape.len() != 3 {
            // Not token embeddings, return as-is
            return data.to_vec();
        }

        let seq_len = shape[1] as usize;
        let dim = shape[2] as usize;
        let mut result = vec![0.0f32; dim];
        let mut count = 0.0f32;

        for i in 0..seq_len {
            if i < attention_mask.len() && attention_mask[i] == 1 {
                for j in 0..dim {
                    result[j] += data[i * dim + j];
                }
                count += 1.0;
            }
        }

        if count > 0.0 {
            for val in &mut result {
                *val /= count;
            }
        }

        // L2 normalize
        let norm: f32 = result.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for val in &mut result {
                *val /= norm;
            }
        }

        result
    }
}

impl Embedder for OnnxEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        let results = self.embed_batch(&[text])?;
        Ok(results.into_iter().next().unwrap_or_default())
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        if texts.is_empty() {
            return Ok(Vec::new());
        }

        let mut all_embeddings = Vec::new();

        // Process one at a time for simplicity
        for text in texts {
            // Tokenize
            let encoding = self
                .tokenizer
                .encode(*text, true)
                .map_err(|e| Error::embedding(format!("Tokenization failed: {}", e)))?;

            let ids = encoding.get_ids();
            let mask = encoding.get_attention_mask();
            let seq_len = ids.len().min(self.max_length);

            // Prepare input tensors
            let input_ids: Vec<i64> = ids.iter().take(seq_len).map(|&id| id as i64).collect();
            let attention_mask: Vec<i64> = mask.iter().take(seq_len).map(|&m| m as i64).collect();
            let token_type_ids: Vec<i64> = vec![0i64; seq_len];

            // Create ort tensors
            let input_ids_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), input_ids.clone())
                .map_err(|e| Error::embedding(format!("Failed to create input tensor: {}", e)))?;
            let input_ids_tensor = Tensor::from_array(input_ids_array)
                .map_err(|e| Error::embedding(format!("Failed to create input tensor: {}", e)))?;

            let attention_mask_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), attention_mask.clone())
                .map_err(|e| Error::embedding(format!("Failed to create mask tensor: {}", e)))?;
            let attention_mask_tensor = Tensor::from_array(attention_mask_array)
                .map_err(|e| Error::embedding(format!("Failed to create mask tensor: {}", e)))?;

            let token_type_ids_array = ndarray::Array::from_shape_vec(IxDyn(&[1, seq_len]), token_type_ids)
                .map_err(|e| Error::embedding(format!("Failed to create token_type tensor: {}", e)))?;
            let token_type_ids_tensor = Tensor::from_array(token_type_ids_array)
                .map_err(|e| Error::embedding(format!("Failed to create token_type tensor: {}", e)))?;

            // Run inference
            let mut session = self.session.lock();
            
            // Get first output name before running
            let first_output_name = session.outputs.first()
                .map(|o| o.name.clone())
                .unwrap_or_else(|| "output".to_string());
            
            let outputs = session
                .run(ort::inputs![
                    "input_ids" => input_ids_tensor,
                    "attention_mask" => attention_mask_tensor,
                    "token_type_ids" => token_type_ids_tensor,
                ])
                .map_err(|e| Error::embedding(format!("Inference failed: {}", e)))?;

            // Extract embeddings - try different output names
            let output = if let Some(val) = outputs.get("last_hidden_state") {
                val
            } else if let Some(val) = outputs.get("sentence_embedding") {
                val
            } else {
                outputs.get(&first_output_name)
                    .ok_or_else(|| Error::embedding("No output found".to_string()))?
            };

            let (output_shape, output_data) = output
                .try_extract_tensor::<f32>()
                .map_err(|e| Error::embedding(format!("Failed to extract output: {}", e)))?;

            let shape_vec: Vec<i64> = output_shape.iter().map(|&d| d as i64).collect();

            let embedding = if shape_vec.len() == 2 {
                // Direct sentence embedding [1, dim]
                let emb: Vec<f32> = output_data.to_vec();
                normalize_vector(emb)
            } else if shape_vec.len() == 3 {
                // Token embeddings [1, seq, dim] - need pooling
                self.mean_pooling(output_data, &shape_vec, &attention_mask)
            } else {
                return Err(Error::embedding(format!(
                    "Unexpected output shape: {:?}",
                    shape_vec
                )));
            };

            all_embeddings.push(embedding);
        }

        Ok(all_embeddings)
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn name(&self) -> &'static str {
        "onnx-runtime"
    }

    fn max_sequence_length(&self) -> usize {
        self.max_length
    }
}

/// Normalize vector to unit length.
fn normalize_vector(mut v: Vec<f32>) -> Vec<f32> {
    let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in &mut v {
            *x /= norm;
        }
    }
    v
}

/// Simple hash-based embedder for testing (no model required).
///
/// Generates deterministic embeddings based on text hash.
/// Useful for testing the search pipeline without loading a real model.
pub struct HashEmbedder {
    dimension: usize,
}

impl HashEmbedder {
    /// Create a new hash embedder with the specified dimension.
    pub fn new(dimension: usize) -> Self {
        Self { dimension }
    }

    fn hash_to_embedding(&self, text: &str) -> Vec<f32> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};

        let mut result = vec![0.0f32; self.dimension];
        let hash = {
            let mut hasher = DefaultHasher::new();
            text.hash(&mut hasher);
            hasher.finish()
        };

        // Generate pseudo-random values from hash
        let mut seed = hash;
        for val in result.iter_mut() {
            seed = seed.wrapping_mul(6364136223846793005).wrapping_add(1);
            *val = ((seed >> 32) as f32 / u32::MAX as f32) * 2.0 - 1.0;
        }

        normalize_vector(result)
    }
}

impl Embedder for HashEmbedder {
    fn embed(&self, text: &str) -> Result<Vec<f32>> {
        Ok(self.hash_to_embedding(text))
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
        Ok(texts.iter().map(|t| self.hash_to_embedding(t)).collect())
    }

    fn dimension(&self) -> usize {
        self.dimension
    }

    fn name(&self) -> &'static str {
        "hash"
    }
}

/// Configuration for loading an embedding model.
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct ModelConfig {
    /// Path to model directory or ONNX file.
    pub model_path: std::path::PathBuf,
    /// Path to tokenizer.json (optional if in model_path).
    pub tokenizer_path: Option<std::path::PathBuf>,
    /// Embedding dimension (auto-detect if None).
    pub dimension: Option<usize>,
    /// Maximum sequence length.
    pub max_length: usize,
}

impl Default for ModelConfig {
    fn default() -> Self {
        Self {
            model_path: std::path::PathBuf::new(),
            tokenizer_path: None,
            dimension: None,
            max_length: 512,
        }
    }
}

impl ModelConfig {
    /// Create config for a model directory.
    pub fn from_directory<P: AsRef<Path>>(path: P) -> Self {
        Self {
            model_path: path.as_ref().to_path_buf(),
            tokenizer_path: None,
            dimension: None,
            max_length: 512,
        }
    }

    /// Set maximum sequence length.
    pub fn with_max_length(mut self, max_length: usize) -> Self {
        self.max_length = max_length;
        self
    }

    /// Set embedding dimension.
    pub fn with_dimension(mut self, dimension: usize) -> Self {
        self.dimension = Some(dimension);
        self
    }

    /// Load the embedder from this config.
    pub fn load(&self) -> Result<OnnxEmbedder> {
        if self.model_path.is_dir() {
            let mut embedder = OnnxEmbedder::from_directory(&self.model_path)?;
            embedder.max_length = self.max_length;
            if let Some(dim) = self.dimension {
                embedder.dimension = dim;
            }
            Ok(embedder)
        } else {
            let tokenizer_path = self
                .tokenizer_path
                .clone()
                .unwrap_or_else(|| self.model_path.with_file_name("tokenizer.json"));

            OnnxEmbedder::new(
                &self.model_path,
                &tokenizer_path,
                self.dimension,
                self.max_length,
            )
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_hash_embedder() {
        let embedder = HashEmbedder::new(384);

        let embedding = embedder.embed("test code").unwrap();
        assert_eq!(embedding.len(), 384);

        // Check normalization
        let norm: f32 = embedding.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 0.01);

        // Same text = same embedding
        let embedding2 = embedder.embed("test code").unwrap();
        assert_eq!(embedding, embedding2);

        // Different text = different embedding
        let embedding3 = embedder.embed("other code").unwrap();
        assert_ne!(embedding, embedding3);
    }

    #[test]
    fn test_batch_embedding() {
        let embedder = HashEmbedder::new(128);

        let texts = vec!["hello", "world", "test"];
        let embeddings = embedder.embed_batch(&texts).unwrap();

        assert_eq!(embeddings.len(), 3);
        for emb in &embeddings {
            assert_eq!(emb.len(), 128);
        }
    }
}