rexis_rag/reranking/

neural_reranker.rs

//! # Neural Reranking
//!
//! Advanced neural network models for reranking including attention mechanisms,
//! transformer architectures, and pre-trained language models.

use crate::{RragResult, SearchResult};
use std::collections::HashMap;

/// Neural reranker with various architecture options
pub struct NeuralReranker {
    /// Configuration
    config: NeuralConfig,

    /// Neural model implementation
    model: Box<dyn NeuralRerankingModel>,

    /// Tokenizer for text preprocessing
    tokenizer: Box<dyn Tokenizer>,

    /// Model cache for performance
    prediction_cache: HashMap<String, f32>,
}

/// Configuration for neural reranking
#[derive(Debug, Clone)]
pub struct NeuralConfig {
    /// Model architecture type
    pub architecture: NeuralArchitecture,

    /// Model parameters
    pub model_params: NeuralModelParams,

    /// Tokenization configuration
    pub tokenization: TokenizationConfig,

    /// Inference configuration
    pub inference_config: InferenceConfig,

    /// Enable prediction caching
    pub enable_caching: bool,

    /// Batch size for inference
    pub batch_size: usize,
}

impl Default for NeuralConfig {
    fn default() -> Self {
        Self {
            architecture: NeuralArchitecture::SimulatedBERT,
            model_params: NeuralModelParams::default(),
            tokenization: TokenizationConfig::default(),
            inference_config: InferenceConfig::default(),
            enable_caching: true,
            batch_size: 16,
        }
    }
}

/// Neural architecture types
#[derive(Debug, Clone, PartialEq)]
pub enum NeuralArchitecture {
    /// BERT-based reranker
    BERT,
    /// RoBERTa-based reranker  
    RoBERTa,
    /// ELECTRA-based reranker
    ELECTRA,
    /// Custom transformer architecture
    CustomTransformer,
    /// Dense neural network
    DenseNetwork,
    /// Convolutional neural network
    CNN,
    /// Recurrent neural network (LSTM/GRU)
    RNN,
    /// Simulated BERT for demonstration
    SimulatedBERT,
}

/// Neural model parameters
#[derive(Debug, Clone)]
pub struct NeuralModelParams {
    /// Hidden dimension size
    pub hidden_dim: usize,

    /// Number of attention heads
    pub num_heads: usize,

    /// Number of layers
    pub num_layers: usize,

    /// Dropout rate
    pub dropout_rate: f32,

    /// Activation function
    pub activation: ActivationFunction,

    /// Maximum sequence length
    pub max_sequence_length: usize,

    /// Model-specific parameters
    pub custom_params: HashMap<String, f32>,
}

impl Default for NeuralModelParams {
    fn default() -> Self {
        Self {
            hidden_dim: 768,
            num_heads: 12,
            num_layers: 12,
            dropout_rate: 0.1,
            activation: ActivationFunction::GELU,
            max_sequence_length: 512,
            custom_params: HashMap::new(),
        }
    }
}

/// Activation functions for neural models
#[derive(Debug, Clone, PartialEq)]
pub enum ActivationFunction {
    ReLU,
    GELU,
    Swish,
    Tanh,
    Sigmoid,
}

/// Tokenization configuration
#[derive(Debug, Clone)]
pub struct TokenizationConfig {
    /// Tokenizer type
    pub tokenizer_type: TokenizerType,

    /// Vocabulary size
    pub vocab_size: usize,

    /// Special tokens
    pub special_tokens: SpecialTokens,

    /// Text preprocessing options
    pub preprocessing: TextPreprocessing,
}

impl Default for TokenizationConfig {
    fn default() -> Self {
        Self {
            tokenizer_type: TokenizerType::WordPiece,
            vocab_size: 30000,
            special_tokens: SpecialTokens::default(),
            preprocessing: TextPreprocessing::default(),
        }
    }
}

/// Types of tokenizers
#[derive(Debug, Clone, PartialEq)]
pub enum TokenizerType {
    WordPiece,
    BPE,
    SentencePiece,
    Whitespace,
    Custom(String),
}

/// Special tokens for neural models
#[derive(Debug, Clone)]
pub struct SpecialTokens {
    /// Classification token
    pub cls_token: String,

    /// Separator token
    pub sep_token: String,

    /// Padding token
    pub pad_token: String,

    /// Unknown token
    pub unk_token: String,

    /// Mask token (for masked language modeling)
    pub mask_token: String,
}

impl Default for SpecialTokens {
    fn default() -> Self {
        Self {
            cls_token: "[CLS]".to_string(),
            sep_token: "[SEP]".to_string(),
            pad_token: "[PAD]".to_string(),
            unk_token: "[UNK]".to_string(),
            mask_token: "[MASK]".to_string(),
        }
    }
}

/// Text preprocessing configuration
#[derive(Debug, Clone)]
pub struct TextPreprocessing {
    /// Convert to lowercase
    pub lowercase: bool,

    /// Remove punctuation
    pub remove_punctuation: bool,

    /// Normalize whitespace
    pub normalize_whitespace: bool,

    /// Remove accents
    pub remove_accents: bool,
}

impl Default for TextPreprocessing {
    fn default() -> Self {
        Self {
            lowercase: true,
            remove_punctuation: false,
            normalize_whitespace: true,
            remove_accents: false,
        }
    }
}

/// Inference configuration
#[derive(Debug, Clone)]
pub struct InferenceConfig {
    /// Use mixed precision (fp16)
    pub use_mixed_precision: bool,

    /// Enable gradient checkpointing (memory optimization)
    pub gradient_checkpointing: bool,

    /// Attention mechanism configuration
    pub attention_config: AttentionConfig,

    /// Output configuration
    pub output_config: OutputConfig,
}

impl Default for InferenceConfig {
    fn default() -> Self {
        Self {
            use_mixed_precision: false,
            gradient_checkpointing: false,
            attention_config: AttentionConfig::default(),
            output_config: OutputConfig::default(),
        }
    }
}

/// Attention mechanism configuration
#[derive(Debug, Clone)]
pub struct AttentionConfig {
    /// Attention mechanism type
    pub mechanism: AttentionMechanism,

    /// Enable attention visualization
    pub enable_visualization: bool,

    /// Attention dropout rate
    pub attention_dropout: f32,

    /// Use relative position encoding
    pub relative_position_encoding: bool,
}

impl Default for AttentionConfig {
    fn default() -> Self {
        Self {
            mechanism: AttentionMechanism::MultiHead,
            enable_visualization: false,
            attention_dropout: 0.1,
            relative_position_encoding: false,
        }
    }
}

/// Types of attention mechanisms
#[derive(Debug, Clone, PartialEq)]
pub enum AttentionMechanism {
    /// Multi-head attention
    MultiHead,
    /// Self-attention
    SelfAttention,
    /// Cross-attention
    CrossAttention,
    /// Sparse attention
    SparseAttention,
    /// Linear attention
    LinearAttention,
}

/// Output configuration for neural models
#[derive(Debug, Clone)]
pub struct OutputConfig {
    /// Output type
    pub output_type: OutputType,

    /// Number of output classes (for classification)
    pub num_classes: Option<usize>,

    /// Enable confidence scores
    pub include_confidence: bool,

    /// Enable attention weights in output
    pub include_attention_weights: bool,
}

impl Default for OutputConfig {
    fn default() -> Self {
        Self {
            output_type: OutputType::RegressionScore,
            num_classes: None,
            include_confidence: true,
            include_attention_weights: false,
        }
    }
}

/// Types of model outputs
#[derive(Debug, Clone, PartialEq)]
pub enum OutputType {
    /// Single relevance score (regression)
    RegressionScore,
    /// Classification probabilities
    Classification,
    /// Ranking scores
    Ranking,
    /// Feature embeddings
    Embeddings,
}

/// Trait for neural reranking models
pub trait NeuralRerankingModel: Send + Sync {
    /// Predict relevance scores for query-document pairs
    fn predict(&self, inputs: &[NeuralInput]) -> RragResult<Vec<NeuralOutput>>;

    /// Predict in batch with specified batch size
    fn predict_batch(
        &self,
        inputs: &[NeuralInput],
        batch_size: usize,
    ) -> RragResult<Vec<NeuralOutput>> {
        let mut results = Vec::new();

        for chunk in inputs.chunks(batch_size) {
            let batch_results = self.predict(chunk)?;
            results.extend(batch_results);
        }

        Ok(results)
    }

    /// Get model information
    fn model_info(&self) -> NeuralModelInfo;

    /// Get attention weights if supported
    fn get_attention_weights(&self, input: &NeuralInput) -> RragResult<Option<AttentionWeights>> {
        let _ = input;
        Ok(None)
    }
}

/// Input to neural reranking models
#[derive(Debug, Clone)]
pub struct NeuralInput {
    /// Query text
    pub query: String,

    /// Document text
    pub document: String,

    /// Tokenized input (if pre-tokenized)
    pub tokens: Option<TokenizedInput>,

    /// Additional features
    pub features: Option<Vec<f32>>,

    /// Input metadata
    pub metadata: NeuralInputMetadata,
}

/// Tokenized input representation
#[derive(Debug, Clone)]
pub struct TokenizedInput {
    /// Token IDs
    pub input_ids: Vec<usize>,

    /// Attention mask
    pub attention_mask: Vec<f32>,

    /// Token type IDs (for BERT-style models)
    pub token_type_ids: Option<Vec<usize>>,

    /// Position IDs
    pub position_ids: Option<Vec<usize>>,
}

/// Metadata for neural input
#[derive(Debug, Clone)]
pub struct NeuralInputMetadata {
    /// Input sequence length
    pub sequence_length: usize,

    /// Number of query tokens
    pub num_query_tokens: usize,

    /// Number of document tokens
    pub num_document_tokens: usize,

    /// Whether input was truncated
    pub truncated: bool,
}

/// Output from neural reranking models
#[derive(Debug, Clone)]
pub struct NeuralOutput {
    /// Relevance score
    pub score: f32,

    /// Confidence in the score
    pub confidence: Option<f32>,

    /// Classification probabilities (if applicable)
    pub probabilities: Option<Vec<f32>>,

    /// Feature embeddings (if requested)
    pub embeddings: Option<Vec<f32>>,

    /// Attention weights (if requested)
    pub attention_weights: Option<AttentionWeights>,

    /// Output metadata
    pub metadata: NeuralOutputMetadata,
}

/// Attention weights from neural models
#[derive(Debug, Clone)]
pub struct AttentionWeights {
    /// Attention weights matrix (layers x heads x seq_len x seq_len)
    pub weights: Vec<Vec<Vec<Vec<f32>>>>,

    /// Token-level attention scores
    pub token_scores: Vec<f32>,

    /// Query-document cross-attention
    pub cross_attention: Option<Vec<Vec<f32>>>,
}

/// Metadata for neural output
#[derive(Debug, Clone)]
pub struct NeuralOutputMetadata {
    /// Model used
    pub model_name: String,

    /// Inference time in milliseconds
    pub inference_time_ms: u64,

    /// Memory usage during inference
    pub memory_usage_mb: Option<f32>,

    /// Model version
    pub model_version: String,
}

/// Information about neural models
#[derive(Debug, Clone)]
pub struct NeuralModelInfo {
    /// Model name
    pub name: String,

    /// Architecture type
    pub architecture: NeuralArchitecture,

    /// Model parameters
    pub parameters: NeuralModelParams,

    /// Number of trainable parameters
    pub num_parameters: Option<usize>,

    /// Model size on disk (MB)
    pub model_size_mb: Option<f32>,

    /// Supported input types
    pub supported_inputs: Vec<String>,

    /// Performance characteristics
    pub performance: ModelPerformance,
}

/// Performance characteristics of neural models
#[derive(Debug, Clone)]
pub struct ModelPerformance {
    /// Average inference time per example (ms)
    pub avg_inference_time_ms: f32,

    /// Memory usage (MB)
    pub memory_usage_mb: f32,

    /// Throughput (examples per second)
    pub throughput: f32,

    /// Accuracy metrics
    pub accuracy_metrics: HashMap<String, f32>,
}

/// Trait for tokenizers
pub trait Tokenizer: Send + Sync {
    /// Tokenize text into tokens
    fn tokenize(&self, text: &str) -> RragResult<Vec<String>>;

    /// Convert tokens to IDs
    fn tokens_to_ids(&self, tokens: &[String]) -> RragResult<Vec<usize>>;

    /// Convert IDs back to tokens
    fn ids_to_tokens(&self, ids: &[usize]) -> RragResult<Vec<String>>;

    /// Tokenize and convert to IDs in one step
    fn encode(&self, text: &str) -> RragResult<Vec<usize>> {
        let tokens = self.tokenize(text)?;
        self.tokens_to_ids(&tokens)
    }

    /// Create tokenized input for query-document pair
    fn create_input(
        &self,
        query: &str,
        document: &str,
        max_length: usize,
    ) -> RragResult<TokenizedInput>;

    /// Get vocabulary size
    fn vocab_size(&self) -> usize;

    /// Get special token IDs
    fn special_tokens(&self) -> &SpecialTokens;
}

impl NeuralReranker {
    /// Create a new neural reranker
    pub fn new(config: NeuralConfig) -> Self {
        let model = Self::create_model(&config);
        let tokenizer = Self::create_tokenizer(&config.tokenization);

        Self {
            config,
            model,
            tokenizer,
            prediction_cache: HashMap::new(),
        }
    }

    /// Create neural model based on configuration
    fn create_model(config: &NeuralConfig) -> Box<dyn NeuralRerankingModel> {
        match &config.architecture {
            NeuralArchitecture::SimulatedBERT => {
                Box::new(SimulatedBertReranker::new(config.model_params.clone()))
            }
            NeuralArchitecture::BERT => Box::new(BertReranker::new(config.model_params.clone())),
            NeuralArchitecture::RoBERTa => {
                Box::new(RobertaReranker::new(config.model_params.clone()))
            }
            _ => {
                // Default to simulated BERT
                Box::new(SimulatedBertReranker::new(config.model_params.clone()))
            }
        }
    }

    /// Create tokenizer based on configuration
    fn create_tokenizer(config: &TokenizationConfig) -> Box<dyn Tokenizer> {
        match config.tokenizer_type {
            TokenizerType::WordPiece => Box::new(SimpleTokenizer::new(config.clone())),
            _ => Box::new(SimpleTokenizer::new(config.clone())),
        }
    }

    /// Rerank search results using neural model
    pub async fn rerank(
        &self,
        query: &str,
        results: &[SearchResult],
    ) -> RragResult<HashMap<usize, f32>> {
        // Create neural inputs
        let inputs: Vec<NeuralInput> = results
            .iter()
            .enumerate()
            .map(|(_idx, result)| {
                let tokenized = self
                    .tokenizer
                    .create_input(
                        query,
                        &result.content,
                        self.config.model_params.max_sequence_length,
                    )
                    .ok();

                NeuralInput {
                    query: query.to_string(),
                    document: result.content.clone(),
                    tokens: tokenized,
                    features: None,
                    metadata: NeuralInputMetadata {
                        sequence_length: query.len() + result.content.len(),
                        num_query_tokens: query.split_whitespace().count(),
                        num_document_tokens: result.content.split_whitespace().count(),
                        truncated: false,
                    },
                }
            })
            .collect();

        // Predict scores
        let outputs = self.model.predict_batch(&inputs, self.config.batch_size)?;

        // Create result mapping
        let mut score_map = HashMap::new();
        for (idx, output) in outputs.into_iter().enumerate() {
            score_map.insert(idx, output.score);
        }

        Ok(score_map)
    }
}

// Convenience type aliases for specific architectures
pub type TransformerReranker = NeuralReranker;
pub type BertReranker = SimulatedBertReranker;
pub type RobertaReranker = SimulatedRobertaReranker;

// Mock implementations
pub struct SimulatedBertReranker {
    params: NeuralModelParams,
}

impl SimulatedBertReranker {
    fn new(params: NeuralModelParams) -> Self {
        Self { params }
    }
}

impl NeuralRerankingModel for SimulatedBertReranker {
    fn predict(&self, inputs: &[NeuralInput]) -> RragResult<Vec<NeuralOutput>> {
        let mut outputs = Vec::new();

        for input in inputs {
            // Simulate BERT-style relevance scoring
            let query_tokens: Vec<&str> = input.query.split_whitespace().collect();
            let doc_tokens: Vec<&str> = input.document.split_whitespace().collect();

            // Simulate attention-based scoring
            let mut attention_score = 0.0;
            let mut total_attention = 0.0;

            for q_token in &query_tokens {
                for d_token in &doc_tokens {
                    let similarity = self.token_similarity(q_token, d_token);
                    let attention_weight = similarity.powf(2.0); // Simulate attention
                    attention_score += similarity * attention_weight;
                    total_attention += attention_weight;
                }
            }

            let normalized_score = if total_attention > 0.0 {
                attention_score / total_attention
            } else {
                0.0
            };

            // Apply sigmoid activation
            let final_score = 1.0 / (1.0 + (-normalized_score * 4.0).exp());

            outputs.push(NeuralOutput {
                score: final_score,
                confidence: Some(0.8),
                probabilities: None,
                embeddings: None,
                attention_weights: None,
                metadata: NeuralOutputMetadata {
                    model_name: "SimulatedBERT".to_string(),
                    inference_time_ms: 10,
                    memory_usage_mb: Some(100.0),
                    model_version: "1.0".to_string(),
                },
            });
        }

        Ok(outputs)
    }

    fn model_info(&self) -> NeuralModelInfo {
        NeuralModelInfo {
            name: "SimulatedBERT-Reranker".to_string(),
            architecture: NeuralArchitecture::SimulatedBERT,
            parameters: self.params.clone(),
            num_parameters: Some(110_000_000),
            model_size_mb: Some(440.0),
            supported_inputs: vec!["text".to_string()],
            performance: ModelPerformance {
                avg_inference_time_ms: 10.0,
                memory_usage_mb: 100.0,
                throughput: 100.0,
                accuracy_metrics: HashMap::new(),
            },
        }
    }
}

impl SimulatedBertReranker {
    fn token_similarity(&self, token1: &str, token2: &str) -> f32 {
        let t1_lower = token1.to_lowercase();
        let t2_lower = token2.to_lowercase();

        if t1_lower == t2_lower {
            1.0
        } else if t1_lower.contains(&t2_lower) || t2_lower.contains(&t1_lower) {
            0.7
        } else {
            // Simple character overlap
            let chars1: std::collections::HashSet<char> = t1_lower.chars().collect();
            let chars2: std::collections::HashSet<char> = t2_lower.chars().collect();

            let intersection = chars1.intersection(&chars2).count();
            let union = chars1.union(&chars2).count();

            if union == 0 {
                0.0
            } else {
                (intersection as f32 / union as f32) * 0.5
            }
        }
    }
}

pub struct SimulatedRobertaReranker {
    params: NeuralModelParams,
}

impl SimulatedRobertaReranker {
    fn new(params: NeuralModelParams) -> Self {
        Self { params }
    }
}

impl NeuralRerankingModel for SimulatedRobertaReranker {
    fn predict(&self, inputs: &[NeuralInput]) -> RragResult<Vec<NeuralOutput>> {
        // Similar to BERT but with slight differences
        let bert_reranker = SimulatedBertReranker::new(self.params.clone());
        let mut outputs = bert_reranker.predict(inputs)?;

        // RoBERTa adjustments
        for output in &mut outputs {
            output.score = (output.score * 1.05).min(1.0); // Slight boost
            output.metadata.model_name = "SimulatedRoBERTa".to_string();
        }

        Ok(outputs)
    }

    fn model_info(&self) -> NeuralModelInfo {
        let mut info = SimulatedBertReranker::new(self.params.clone()).model_info();
        info.name = "SimulatedRoBERTa-Reranker".to_string();
        info.architecture = NeuralArchitecture::RoBERTa;
        info.num_parameters = Some(125_000_000);
        info
    }
}

// Simple tokenizer implementation
struct SimpleTokenizer {
    config: TokenizationConfig,
}

impl SimpleTokenizer {
    fn new(config: TokenizationConfig) -> Self {
        Self { config }
    }
}

impl Tokenizer for SimpleTokenizer {
    fn tokenize(&self, text: &str) -> RragResult<Vec<String>> {
        let mut processed_text = text.to_string();

        if self.config.preprocessing.lowercase {
            processed_text = processed_text.to_lowercase();
        }

        if self.config.preprocessing.normalize_whitespace {
            processed_text = processed_text
                .split_whitespace()
                .collect::<Vec<_>>()
                .join(" ");
        }

        let tokens: Vec<String> = processed_text
            .split_whitespace()
            .map(|s| s.to_string())
            .collect();

        Ok(tokens)
    }

    fn tokens_to_ids(&self, tokens: &[String]) -> RragResult<Vec<usize>> {
        // Simple hash-based ID assignment
        let ids = tokens
            .iter()
            .map(|token| {
                use std::collections::hash_map::DefaultHasher;
                use std::hash::{Hash, Hasher};

                let mut hasher = DefaultHasher::new();
                token.hash(&mut hasher);
                (hasher.finish() % self.config.vocab_size as u64) as usize
            })
            .collect();

        Ok(ids)
    }

    fn ids_to_tokens(&self, ids: &[usize]) -> RragResult<Vec<String>> {
        // Simple reverse mapping (not accurate for real tokenizers)
        let tokens = ids.iter().map(|&id| format!("token_{}", id)).collect();

        Ok(tokens)
    }

    fn create_input(
        &self,
        query: &str,
        document: &str,
        max_length: usize,
    ) -> RragResult<TokenizedInput> {
        let query_tokens = self.tokenize(query)?;
        let document_tokens = self.tokenize(document)?;

        // Create BERT-style input: [CLS] query [SEP] document [SEP]
        let mut all_tokens = vec![self.config.special_tokens.cls_token.clone()];
        all_tokens.extend(query_tokens);
        all_tokens.push(self.config.special_tokens.sep_token.clone());
        all_tokens.extend(document_tokens);
        all_tokens.push(self.config.special_tokens.sep_token.clone());

        // Truncate if necessary
        if all_tokens.len() > max_length {
            all_tokens.truncate(max_length - 1);
            all_tokens.push(self.config.special_tokens.sep_token.clone());
        }

        // Pad to max_length
        while all_tokens.len() < max_length {
            all_tokens.push(self.config.special_tokens.pad_token.clone());
        }

        let input_ids = self.tokens_to_ids(&all_tokens)?;
        let attention_mask: Vec<f32> = all_tokens
            .iter()
            .map(|token| {
                if token == &self.config.special_tokens.pad_token {
                    0.0
                } else {
                    1.0
                }
            })
            .collect();

        Ok(TokenizedInput {
            input_ids,
            attention_mask,
            token_type_ids: None,
            position_ids: None,
        })
    }

    fn vocab_size(&self) -> usize {
        self.config.vocab_size
    }

    fn special_tokens(&self) -> &SpecialTokens {
        &self.config.special_tokens
    }
}

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

    #[tokio::test]
    async fn test_neural_reranking() {
        let config = NeuralConfig::default();
        let reranker = NeuralReranker::new(config);

        let results = vec![
            SearchResult {
                id: "doc1".to_string(),
                content: "Machine learning algorithms for data analysis".to_string(),
                score: 0.8,
                rank: 0,
                metadata: HashMap::new(),
                embedding: None,
            },
            SearchResult {
                id: "doc2".to_string(),
                content: "Cooking recipes for beginners".to_string(),
                score: 0.3,
                rank: 1,
                metadata: HashMap::new(),
                embedding: None,
            },
        ];

        let query = "machine learning data science";
        let reranked_scores = reranker.rerank(query, &results).await.unwrap();

        assert!(!reranked_scores.is_empty());
        // First document should have higher neural score
        assert!(reranked_scores.get(&0).unwrap() > reranked_scores.get(&1).unwrap());
    }

    #[test]
    fn test_tokenizer() {
        let config = TokenizationConfig::default();
        let tokenizer = SimpleTokenizer::new(config);

        let tokens = tokenizer.tokenize("Hello world!").unwrap();
        assert!(!tokens.is_empty());

        let input = tokenizer.create_input("query", "document", 128).unwrap();
        assert_eq!(input.input_ids.len(), 128);
        assert_eq!(input.attention_mask.len(), 128);
    }

    #[test]
    fn test_simulated_bert() {
        let params = NeuralModelParams::default();
        let model = SimulatedBertReranker::new(params);

        let input = NeuralInput {
            query: "machine learning".to_string(),
            document: "artificial intelligence and machine learning".to_string(),
            tokens: None,
            features: None,
            metadata: NeuralInputMetadata {
                sequence_length: 50,
                num_query_tokens: 2,
                num_document_tokens: 5,
                truncated: false,
            },
        };

        let outputs = model.predict(&[input]).unwrap();
        assert_eq!(outputs.len(), 1);
        assert!(outputs[0].score >= 0.0 && outputs[0].score <= 1.0);
        assert!(outputs[0].confidence.is_some());
    }
}