fix skipped tests (#621)

Signed-off-by: Yehudit Kerido <[email protected]>
Co-authored-by: Yehudit Kerido <[email protected]>

Author: yehudit1987

candle-binding/semantic-router_test.go

@@ -1362,15 +1362,15 @@ func TestCandleBertTokensWithLabels(t *testing.T) {
 
 	success := InitCandleBertTokenClassifier(BertPIITokenClassifierModelPath, 9, true) // 9 PII classes
 	if !success {
-		t.Skipf("Candle BERT token classifier not available at path: %s", BertPIITokenClassifierModelPath)
+		t.Fatalf("Failed to initialize Candle BERT token classifier at path: %s. Model should be available in CI (included in LoRA model set).", BertPIITokenClassifierModelPath)
 	}
 
 	testText := "Contact Dr. Sarah Johnson at [email protected] for medical records"
 
 	result, err := ClassifyCandleBertTokensWithLabels(testText, id2labelJSON)
 	if err != nil {
 		if isModelInitializationError(err) {
-			t.Skipf("Skipping Candle BERT token classifier tests due to model initialization error: %v", err)
+			t.Fatalf("Candle BERT token classifier failed with model initialization error: %v. Model should be initialized earlier in test.", err)
 		}
 		t.Fatalf("Token classification with labels failed: %v", err)
 	}

candle-binding/src/classifiers/lora/token_lora.rs

@@ -187,20 +187,23 @@ impl LoRATokenClassifier {
         for (i, (token, token_embedding)) in tokens.iter().enumerate() {
             // Use real BERT embedding from tokenization
 
+            // Add batch dimension: [hidden_size] -> [1, hidden_size]
+            let token_embedding_batched = token_embedding.unsqueeze(0)?;
+
             // Apply base classifier
-            let base_logits = self.base_classifier.forward(&token_embedding)?;
+            let base_logits = self.base_classifier.forward(&token_embedding_batched)?;
 
             // Apply LoRA adapters if available
             let enhanced_logits = if let Some(adapter) = self.adapters.get("token_classification") {
-                let adapter_output = adapter.forward(&token_embedding, false)?; // false = not training
+                let adapter_output = adapter.forward(&token_embedding_batched, false)?; // false = not training
                 (&base_logits + &adapter_output)?
             } else {
                 base_logits
             };
 
-            // Apply softmax to get probabilities
+            // Apply softmax to get probabilities and remove batch dimension
             let probabilities = candle_nn::ops::softmax(&enhanced_logits, 1)?;
-            let probs_vec = probabilities.to_vec1::<f32>()?;
+            let probs_vec = probabilities.squeeze(0)?.to_vec1::<f32>()?;
 
             // Find the class with highest probability
             let (predicted_id, confidence) = probs_vec

candle-binding/src/ffi/classify.rs

@@ -520,8 +520,40 @@ pub extern "C" fn classify_candle_bert_tokens_with_labels(
         }
     };
 
-    // Use TraditionalBertTokenClassifier for token-level classification with labels
+    // Intelligent routing: Check LoRA token classifier first, then fall back to traditional
 
+    // Try LoRA token classifier first
+    if let Some(classifier) = crate::ffi::init::LORA_TOKEN_CLASSIFIER.get() {
+        let classifier = classifier.clone();
+        match classifier.classify_tokens(text) {
+            Ok(token_results) => {
+                // Filter out "O" (Outside) labels - only return actual entities
+                let token_entities: Vec<(String, String, f32)> = token_results
+                    .iter()
+                    .filter(|result| result.label_name != "O" && result.label_id != 0)
+                    .map(|result| {
+                        (
+                            result.token.clone(),
+                            result.label_name.clone(),
+                            result.confidence,
+                        )
+                    })
+                    .collect();
+
+                let entities_ptr = unsafe { allocate_bert_token_entity_array(&token_entities) };
+
+                return BertTokenClassificationResult {
+                    entities: entities_ptr,
+                    num_entities: token_entities.len() as i32,
+                };
+            }
+            Err(_e) => {
+                // Fall through to traditional classifier
+            }
+        }
+    }
+
+    // Fall back to traditional BERT token classifier
     if let Some(classifier) = TRADITIONAL_BERT_TOKEN_CLASSIFIER.get() {
         let classifier = classifier.clone();
         match classifier.classify_tokens(text) {
@@ -581,17 +613,19 @@ pub extern "C" fn classify_candle_bert_tokens(
         let lora_classifier = lora_classifier.clone();
         match lora_classifier.classify_tokens(text) {
             Ok(lora_results) => {
+                // Filter out "O" (Outside) labels - only return actual entities
                 // Convert LoRA results to BertTokenEntity format
                 let token_entities: Vec<(String, String, f32)> = lora_results
                     .iter()
+                    .filter(|r| r.label_name != "O" && r.label_id != 0)
                     .map(|r| (r.token.clone(), r.label_name.clone(), r.confidence))
                     .collect();
 
                 let entities_ptr = unsafe { allocate_bert_token_entity_array(&token_entities) };
 
                 return BertTokenClassificationResult {
                     entities: entities_ptr,
-                    num_entities: lora_results.len() as i32,
+                    num_entities: token_entities.len() as i32,
                 };
             }
             Err(_e) => {

candle-binding/src/ffi/init.rs

@@ -594,6 +594,11 @@ pub extern "C" fn init_candle_bert_token_classifier(
 
     match model_type {
         ModelType::LoRA => {
+            // Check if already initialized
+            if LORA_TOKEN_CLASSIFIER.get().is_some() {
+                return true; // Already initialized, return success
+            }
+
             // Route to LoRA token classifier initialization
             match crate::classifiers::lora::token_lora::LoRATokenClassifier::new(
                 model_path, use_cpu,
@@ -606,6 +611,14 @@ pub extern "C" fn init_candle_bert_token_classifier(
             }
         }
         ModelType::Traditional => {
+            // Check if already initialized
+            if crate::model_architectures::traditional::bert::TRADITIONAL_BERT_TOKEN_CLASSIFIER
+                .get()
+                .is_some()
+            {
+                return true; // Already initialized, return success
+            }
+
             // Route to traditional BERT token classifier
             match crate::model_architectures::traditional::bert::TraditionalBertTokenClassifier::new(
                 model_path,

tools/make/rust.mk

@@ -64,8 +64,7 @@ test-binding-lora: $(if $(CI),rust-ci,rust) ## Run Go tests with LoRA and advanc
 	@echo "Running candle-binding tests with LoRA and advanced embedding models..."
 	@export LD_LIBRARY_PATH=${PWD}/candle-binding/target/release && \
 		cd candle-binding && CGO_ENABLED=1 go test -v -race \
-		-run "^Test(BertTokenClassification|BertSequenceClassification|CandleBertClassifier|CandleBertTokenClassifier|CandleBertTokensWithLabels|LoRAUnifiedClassifier|GetEmbeddingSmart|InitEmbeddingModels|GetEmbeddingWithDim|EmbeddingConsistency|EmbeddingPriorityRouting|EmbeddingConcurrency)$$" \
-		|| { echo "⚠️  Warning: Some LoRA/embedding tests failed (may be due to missing restricted models), continuing..."; $(if $(CI),true,exit 1); }
+		-run "^Test(BertTokenClassification|BertSequenceClassification|CandleBertClassifier|CandleBertTokenClassifier|CandleBertTokensWithLabels|LoRAUnifiedClassifier|GetEmbeddingSmart|InitEmbeddingModels|GetEmbeddingWithDim|EmbeddingConsistency|EmbeddingPriorityRouting|EmbeddingConcurrency)$$"
 
 # Test the Rust library - all tests (conditionally use rust-ci in CI environments)
 test-binding: $(if $(CI),rust-ci,rust) ## Run all Go tests with the Rust static library