candle_semantic_router/
lib.rs

1// This file is a binding for the candle-core and candle-transformers libraries.
2// It is based on https://github.com/huggingface/candle/tree/main/candle-examples/examples/bert
3use std::ffi::{c_char, CStr, CString};
4use std::sync::Arc;
5use std::sync::Mutex;
6use std::path::Path;
7
8pub mod modernbert;
9
10// Re-export ModernBERT functions and structures
11pub use modernbert::{
12    ModernBertClassificationResult,
13    init_modernbert_classifier,
14    init_modernbert_pii_classifier,
15    init_modernbert_jailbreak_classifier,
16    classify_modernbert_text,
17    classify_modernbert_pii_text,
18    classify_modernbert_jailbreak_text,
19};
20
21use anyhow::{Error as E, Result};
22use candle_core::{DType, Device, Tensor};
23use candle_nn::{VarBuilder, Linear};
24use candle_transformers::models::bert::{BertModel, Config, HiddenAct, DTYPE};
25use hf_hub::{api::sync::Api, Repo, RepoType};
26use tokenizers::Tokenizer;
27use tokenizers::TruncationParams;
28use tokenizers::TruncationStrategy;
29use tokenizers::TruncationDirection;
30
31// Structure to hold BERT model and tokenizer for semantic similarity
32pub struct BertSimilarity {
33    model: BertModel,
34    tokenizer: Tokenizer,
35    device: Device,
36}
37
38// Structure to hold BERT model, tokenizer, and classification head for text classification
39pub struct BertClassifier {
40    model: BertModel,
41    tokenizer: Tokenizer,
42    classification_head: Linear,
43    num_classes: usize,
44    device: Device,
45}
46
47lazy_static::lazy_static! {
48    static ref BERT_SIMILARITY: Arc<Mutex<Option<BertSimilarity>>> = Arc::new(Mutex::new(None));
49    static ref BERT_CLASSIFIER: Arc<Mutex<Option<BertClassifier>>> = Arc::new(Mutex::new(None));
50    static ref BERT_PII_CLASSIFIER: Arc<Mutex<Option<BertClassifier>>> = Arc::new(Mutex::new(None));
51    static ref BERT_JAILBREAK_CLASSIFIER: Arc<Mutex<Option<BertClassifier>>> = Arc::new(Mutex::new(None));
52}
53
54// Structure to hold tokenization result
55#[repr(C)]
56pub struct TokenizationResult {
57    pub token_ids: *mut i32,
58    pub token_count: i32,
59    pub tokens: *mut *mut c_char,
60    pub error: bool,
61}
62
63impl BertSimilarity {
64    pub fn new(model_id: &str, use_cpu: bool) -> Result<Self> {
65        let device = if use_cpu {
66            Device::Cpu
67        } else {
68            Device::cuda_if_available(0)?
69        };
70
71        // Default to a sentence transformer model if not specified or empty
72        let model_id = if model_id.is_empty() {
73            "sentence-transformers/all-MiniLM-L6-v2"
74        } else {
75            model_id
76        };
77
78        let (config_filename, tokenizer_filename, weights_filename, use_pth) = if Path::new(model_id).exists() {
79            // Local model path
80            println!("Loading model from local directory: {}", model_id);
81            let config_path = Path::new(model_id).join("config.json");
82            let tokenizer_path = Path::new(model_id).join("tokenizer.json");
83            
84            // Check for safetensors first, fall back to PyTorch
85            let weights_path = if Path::new(model_id).join("model.safetensors").exists() {
86                (Path::new(model_id).join("model.safetensors").to_string_lossy().to_string(), false)
87            } else if Path::new(model_id).join("pytorch_model.bin").exists() {
88                (Path::new(model_id).join("pytorch_model.bin").to_string_lossy().to_string(), true)
89            } else {
90                return Err(E::msg(format!("No model weights found in {}", model_id)));
91            };
92            
93            (
94                config_path.to_string_lossy().to_string(),
95                tokenizer_path.to_string_lossy().to_string(),
96                weights_path.0,
97                weights_path.1
98            )
99        } else {
100            // HuggingFace Hub model
101            println!("Loading model from HuggingFace Hub: {}", model_id);
102            let repo = Repo::with_revision(
103                model_id.to_string(), 
104                RepoType::Model, 
105                "main".to_string()
106            );
107
108            let api = Api::new()?;
109            let api = api.repo(repo);
110            let config = api.get("config.json")?;
111            let tokenizer = api.get("tokenizer.json")?;
112
113            // Try to get safetensors first, if that fails, fall back to pytorch_model.bin. This is for BAAI models
114            // create a special case for BAAI to download the correct weights to avoid downloading the wrong weights
115            let (weights, use_pth) = if model_id.starts_with("BAAI/") {
116                // BAAI models typically use PyTorch model format
117                (api.get("pytorch_model.bin")?, true)
118            } else {
119                match api.get("model.safetensors") {
120                    Ok(weights) => (weights, false),
121                    Err(_) => {
122                        println!("Safetensors model not found, trying PyTorch model instead...");
123                        (api.get("pytorch_model.bin")?, true)
124                    }
125                }
126            };
127
128            (
129                config.to_string_lossy().to_string(),
130                tokenizer.to_string_lossy().to_string(),
131                weights.to_string_lossy().to_string(),
132                use_pth
133            )
134        };
135
136        let config = std::fs::read_to_string(config_filename)?;
137        let mut config: Config = serde_json::from_str(&config)?;
138        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
139
140        // Use the approximate GELU for better performance
141        config.hidden_act = HiddenAct::GeluApproximate;
142
143        let vb = if use_pth {
144            VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
145        } else {
146            unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
147        };
148
149        let model = BertModel::load(vb, &config)?;
150
151        Ok(Self {
152            model,
153            tokenizer,
154            device,
155        })
156    }
157
158    // Tokenize a text string
159    pub fn tokenize_text(&self, text: &str, max_length: Option<usize>) -> Result<(Vec<i32>, Vec<String>)> {
160        // Encode the text with the tokenizer
161        let mut tokenizer = self.tokenizer.clone();
162        tokenizer.with_truncation(Some(TruncationParams {
163            max_length: max_length.unwrap_or(512),
164            strategy: TruncationStrategy::LongestFirst,
165            stride: 0,
166            direction: TruncationDirection::Right,
167        })).map_err(E::msg)?;
168        
169        let encoding = tokenizer.encode(text, true)
170            .map_err(E::msg)?;
171        
172        // Get token IDs and tokens
173        let token_ids = encoding.get_ids().iter().map(|&id| id as i32).collect();
174        let tokens = encoding.get_tokens().to_vec();
175        
176        Ok((token_ids, tokens))
177    }
178
179    // Get embedding for a text
180    pub fn get_embedding(&self, text: &str, max_length: Option<usize>) -> Result<Tensor> {
181        // Encode the text with the tokenizer
182        let mut tokenizer = self.tokenizer.clone();
183        tokenizer.with_truncation(Some(TruncationParams {
184            max_length: max_length.unwrap_or(512),
185            strategy: TruncationStrategy::LongestFirst,
186            stride: 0,
187            direction: TruncationDirection::Right,
188        })).map_err(E::msg)?;
189        
190        let encoding = tokenizer.encode(text, true)
191            .map_err(E::msg)?;
192        
193        // Get token IDs and attention mask
194        let token_ids = encoding.get_ids().to_vec();
195        let attention_mask = encoding.get_attention_mask().to_vec();
196        
197        // Create tensors
198        let token_ids_tensor = Tensor::new(&token_ids[..], &self.device)?.unsqueeze(0)?;
199        let attention_mask_tensor = Tensor::new(&attention_mask[..], &self.device)?.unsqueeze(0)?;
200        let token_type_ids = token_ids_tensor.zeros_like()?;
201        
202        // Run the text through BERT with attention mask
203        let embeddings = self.model.forward(&token_ids_tensor, &token_type_ids, Some(&attention_mask_tensor))?;
204        
205        // Mean pooling: sum over tokens and divide by attention mask sum
206        let sum_embeddings = embeddings.sum(1)?;
207        let attention_sum = attention_mask_tensor.sum(1)?.to_dtype(embeddings.dtype())?;
208        let pooled = sum_embeddings.broadcast_div(&attention_sum)?;
209        
210        // Convert to float32 and normalize
211        let embedding = pooled.to_dtype(DType::F32)?;
212        
213        normalize_l2(&embedding)
214    }
215
216    // Calculate cosine similarity between two texts
217    pub fn calculate_similarity(&self, text1: &str, text2: &str, max_length: Option<usize>) -> Result<f32> {
218        let embedding1 = self.get_embedding(text1, max_length)?;
219        let embedding2 = self.get_embedding(text2, max_length)?;
220        
221        // For normalized vectors, dot product equals cosine similarity
222        let dot_product = embedding1.matmul(&embedding2.transpose(0, 1)?)?;
223        
224        // Extract the scalar value from the result
225        let sim_value = dot_product.squeeze(0)?.squeeze(0)?.to_scalar::<f32>()?;
226        
227        Ok(sim_value)
228    }
229
230    // Find most similar text from a list
231    pub fn find_most_similar(&self, query_text: &str, candidates: &[&str], max_length: Option<usize>) -> Result<(usize, f32)> {
232        if candidates.is_empty() {
233            return Err(E::msg("Empty candidate list"));
234        }
235        
236        let query_embedding = self.get_embedding(query_text, max_length)?;
237        
238        // Calculate similarity for each candidate individually
239        let mut best_idx = 0;
240        let mut best_score = -1.0;
241        
242        for (idx, candidate) in candidates.iter().enumerate() {
243            let candidate_embedding = self.get_embedding(candidate, max_length)?;
244            
245            // Calculate similarity (dot product of normalized vectors = cosine similarity)
246            let sim = query_embedding.matmul(&candidate_embedding.transpose(0, 1)?)?;
247            let score = sim.squeeze(0)?.squeeze(0)?.to_scalar::<f32>()?;
248            
249            if score > best_score {
250                best_score = score;
251                best_idx = idx;
252            }
253        }
254        
255        Ok((best_idx, best_score))
256    }
257}
258
259impl BertClassifier {
260    pub fn new(model_id: &str, num_classes: usize, use_cpu: bool) -> Result<Self> {
261        if num_classes < 2 {
262            return Err(E::msg(format!("Number of classes must be at least 2, got {}", num_classes)));
263        }
264
265        let device = if use_cpu {
266            Device::Cpu
267        } else {
268            Device::cuda_if_available(0)?
269        };
270
271        println!("Initializing classifier model: {}", model_id);
272
273        // Check if this is a SentenceTransformer linear classifier model
274        let is_sentence_transformer = Path::new(model_id).join("modules.json").exists();
275        
276        if is_sentence_transformer {
277            println!("Detected SentenceTransformer model with linear classifier head");
278        }
279
280        let (config_filename, tokenizer_filename, weights_filename, use_pth) = if Path::new(model_id).exists() {
281            // Local model path
282            println!("Loading model from local directory: {}", model_id);
283            let config_path = Path::new(model_id).join("config.json");
284            let tokenizer_path = Path::new(model_id).join("tokenizer.json");
285            
286            // For SentenceTransformer models, check both the root and 0_Transformer
287            let weights_path = if is_sentence_transformer {
288                // First check if model weights are at the root level (most common for sentence-transformers)
289                if Path::new(model_id).join("model.safetensors").exists() {
290                    println!("Found model weights at root level");
291                    (Path::new(model_id).join("model.safetensors").to_string_lossy().to_string(), false)
292                } else if Path::new(model_id).join("pytorch_model.bin").exists() {
293                    println!("Found PyTorch model at root level");
294                    (Path::new(model_id).join("pytorch_model.bin").to_string_lossy().to_string(), true)
295                }
296                // Otherwise check if there's a 0_Transformer directory
297                else {
298                    let transformer_path = Path::new(model_id).join("0_Transformer");
299                    if transformer_path.exists() {
300                        if transformer_path.join("model.safetensors").exists() {
301                            (transformer_path.join("model.safetensors").to_string_lossy().to_string(), false)
302                        } else if transformer_path.join("pytorch_model.bin").exists() {
303                            (transformer_path.join("pytorch_model.bin").to_string_lossy().to_string(), true)
304                        } else {
305                            return Err(E::msg(format!("No transformer model weights found in {}", transformer_path.display())));
306                        }
307                    } else {
308                        return Err(E::msg(format!("No model weights found in {}", model_id)));
309                    }
310                }
311            } else if Path::new(model_id).join("model.safetensors").exists() {
312                (Path::new(model_id).join("model.safetensors").to_string_lossy().to_string(), false)
313            } else if Path::new(model_id).join("pytorch_model.bin").exists() {
314                (Path::new(model_id).join("pytorch_model.bin").to_string_lossy().to_string(), true)
315            } else {
316                return Err(E::msg(format!("No model weights found in {}", model_id)));
317            };
318            
319            (
320                config_path.to_string_lossy().to_string(),
321                tokenizer_path.to_string_lossy().to_string(),
322                weights_path.0,
323                weights_path.1
324            )
325        } else {
326            // HuggingFace Hub model
327            println!("Loading model from HuggingFace Hub: {}", model_id);
328            let repo = Repo::with_revision(
329                model_id.to_string(),
330                RepoType::Model,
331                "main".to_string(),
332            );
333
334            let api = Api::new()?;
335            let api = api.repo(repo);
336            let config = api.get("config.json")?;
337            let tokenizer = api.get("tokenizer.json")?;
338
339            // Try safetensors first, fall back to PyTorch
340            let (weights, use_pth) = match api.get("model.safetensors") {
341                Ok(weights) => (weights, false),
342                Err(_) => {
343                    println!("Safetensors model not found, trying PyTorch model instead...");
344                    (api.get("pytorch_model.bin")?, true)
345                }
346            };
347
348            (
349                config.to_string_lossy().to_string(),
350                tokenizer.to_string_lossy().to_string(),
351                weights.to_string_lossy().to_string(),
352                use_pth
353            )
354        };
355
356        let config = std::fs::read_to_string(config_filename)?;
357        let mut config: Config = serde_json::from_str(&config)?;
358        let tokenizer = Tokenizer::from_file(tokenizer_filename).map_err(E::msg)?;
359
360        // Use approximate GELU for better performance
361        config.hidden_act = HiddenAct::GeluApproximate;
362
363        let vb = if use_pth {
364            VarBuilder::from_pth(&weights_filename, DTYPE, &device)?
365        } else {
366            unsafe { VarBuilder::from_mmaped_safetensors(&[weights_filename], DTYPE, &device)? }
367        };
368
369        println!("Successfully loaded transformer model");
370        let model = BertModel::load(vb.clone(), &config)?;
371        println!("Successfully initialized BERT model instance");
372
373        // Create a classification head
374        // For SentenceTransformer models, we need to load the Dense layer weights from 2_Dense
375        let (w, b) = if is_sentence_transformer {
376            // Load the dense layer weights from 2_Dense
377            let dense_dir = Path::new(model_id).join("2_Dense");
378            println!("Looking for dense weights in {}", dense_dir.display());
379            
380            let dense_config_path = dense_dir.join("config.json");
381            
382            if dense_config_path.exists() {
383                println!("Found dense config at {}", dense_config_path.display());
384                let dense_config = std::fs::read_to_string(dense_config_path)?;
385                let dense_config: serde_json::Value = serde_json::from_str(&dense_config)?;
386                
387                // Get dimensions from the config
388                let in_features = dense_config["in_features"].as_i64().unwrap_or(768) as usize;
389                let out_features = dense_config["out_features"].as_i64().unwrap_or(num_classes as i64) as usize;
390                
391                println!("Dense layer dimensions: in_features={}, out_features={}", in_features, out_features);
392                
393                // Try to load dense weights from safetensors or pytorch files
394                let weights_path = if dense_dir.join("model.safetensors").exists() {
395                    println!("Found dense safetensors weights");
396                    (dense_dir.join("model.safetensors").to_string_lossy().to_string(), false)
397                } else if dense_dir.join("pytorch_model.bin").exists() {
398                    println!("Found dense PyTorch weights");
399                    (dense_dir.join("pytorch_model.bin").to_string_lossy().to_string(), true)
400                } else {
401                    return Err(E::msg(format!("No dense layer weights found in {}", dense_dir.display())));
402                };
403                
404                // Load the weights
405                let dense_vb = if weights_path.1 {
406                    VarBuilder::from_pth(&weights_path.0, DType::F32, &device)?
407                } else {
408                    unsafe { VarBuilder::from_mmaped_safetensors(&[weights_path.0], DType::F32, &device)? }
409                };
410                
411                // Get the weight and bias tensors - PyTorch uses [out_features, in_features] format
412                let weight = dense_vb.get((out_features, in_features), "linear.weight")?;
413                // Transpose the weight matrix to match our expected format [in_features, out_features]
414                let weight = weight.t()?;
415                let bias = dense_vb.get(out_features, "linear.bias")?;
416                println!("Successfully loaded dense layer weights");
417                
418                (weight, bias)
419            } else {
420                // Fallback: create random weights as before
421                println!("No dense config found, using random weights");
422                let hidden_size = config.hidden_size;
423                let w = Tensor::randn(0.0, 0.02, (hidden_size, num_classes), &device)?;
424                let b = Tensor::zeros((num_classes,), DType::F32, &device)?;
425                (w, b)
426            }
427        } else {
428            // Regular BERT model: create random weights
429            let hidden_size = config.hidden_size;
430            let w = Tensor::randn(0.0, 0.02, (hidden_size, num_classes), &device)?;
431            let b = Tensor::zeros((num_classes,), DType::F32, &device)?;
432            (w, b)
433        };
434        
435        let classification_head = Linear::new(w, Some(b));
436        println!("Linear classification head created");
437
438        Ok(Self {
439            model,
440            tokenizer,
441            classification_head,
442            num_classes,
443            device,
444        })
445    }
446
447    pub fn classify_text(&self, text: &str) -> Result<(usize, f32)> {
448        // Encode the text with the tokenizer
449        let encoding = self.tokenizer
450            .encode(text, true)
451            .map_err(E::msg)?;
452        
453        let token_ids = encoding.get_ids().to_vec();
454        let attention_mask = encoding.get_attention_mask().to_vec();
455        let token_ids_tensor = Tensor::new(&token_ids[..], &self.device)?.unsqueeze(0)?;
456        let token_type_ids = token_ids_tensor.zeros_like()?;
457        let attention_mask_tensor = Tensor::new(&attention_mask[..], &self.device)?.unsqueeze(0)?;
458        
459        // Run the text through BERT
460        let embeddings = self.model.forward(&token_ids_tensor, &token_type_ids, Some(&attention_mask_tensor))?;
461        
462        // Implement proper mean pooling for SentenceTransformer
463        // Sum over token dimension (dim=1) and divide by attention mask sum to get mean
464        let embedding_sum = embeddings.sum(1)?;
465        let attention_mask_sum = attention_mask_tensor.to_dtype(embeddings.dtype())?.sum(1)?;
466        let pooled_embedding = embedding_sum.broadcast_div(&attention_mask_sum)?;
467        
468        // Get the dimensions and convert to the right type
469        let pooled_embedding = pooled_embedding.to_dtype(DType::F32)?;
470        
471        // Apply the linear layer (classification head) manually
472        let weights = self.classification_head.weight().to_dtype(DType::F32)?;
473        let bias = self.classification_head.bias().unwrap().to_dtype(DType::F32)?;
474        
475        // Use matmul with the weights matrix
476        // If weights are already transposed to [in_features, out_features]
477        let logits = pooled_embedding.matmul(&weights)?;
478        
479        // Add bias
480        let logits = logits.broadcast_add(&bias)?;
481        
482        // If logits has shape [1, num_classes], squeeze it to get [num_classes]
483        let logits = if logits.dims().len() > 1 {
484            logits.squeeze(0)?
485        } else {
486            logits
487        };
488        
489        // Apply softmax to get probabilities
490        let logits_vec = logits.to_vec1::<f32>()?;
491        let max_logit = logits_vec.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b));
492        let exp_values: Vec<f32> = logits_vec.iter().map(|&x| (x - max_logit).exp()).collect();
493        let exp_sum: f32 = exp_values.iter().sum();
494        let probabilities: Vec<f32> = exp_values.iter().map(|&x| x / exp_sum).collect();
495        
496        // Get the predicted class with highest probability
497        let (predicted_idx, &max_prob) = probabilities.iter()
498            .enumerate()
499            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
500            .unwrap_or((0, &0.0));
501        
502        // Ensure we don't return a class index outside our expected range
503        if predicted_idx >= self.num_classes {
504            return Err(E::msg(format!(
505                "Invalid class index: {} (num_classes: {})",
506                predicted_idx, self.num_classes
507            )));
508        }
509        
510        Ok((predicted_idx, max_prob))
511    }
512}
513
514// Tokenize text (called from Go)
515#[no_mangle]
516pub extern "C" fn tokenize_text(text: *const c_char, max_length: i32) -> TokenizationResult {
517    let text = unsafe {
518        match CStr::from_ptr(text).to_str() {
519            Ok(s) => s,
520            Err(_) => return TokenizationResult {
521                token_ids: std::ptr::null_mut(),
522                token_count: 0,
523                tokens: std::ptr::null_mut(),
524                error: true
525            },
526        }
527    };
528
529    let bert_opt = BERT_SIMILARITY.lock().unwrap();
530    let bert = match &*bert_opt {
531        Some(b) => b,
532        None => {
533            eprintln!("BERT model not initialized");
534            return TokenizationResult {
535                token_ids: std::ptr::null_mut(),
536                token_count: 0,
537                tokens: std::ptr::null_mut(),
538                error: true
539            };
540        }
541    };
542
543    let max_length_opt = if max_length <= 0 { None } else { Some(max_length as usize) };
544    match bert.tokenize_text(text, max_length_opt) {
545        Ok((token_ids, tokens)) => {
546            let count = token_ids.len() as i32;
547            
548            // Allocate memory for token IDs
549            let ids_ptr = token_ids.as_ptr() as *mut i32;
550            
551            // Allocate memory for tokens
552            let c_tokens: Vec<*mut c_char> = tokens.iter()
553                .map(|s| CString::new(s.as_str()).unwrap().into_raw())
554                .collect();
555            
556            let tokens_ptr = c_tokens.as_ptr() as *mut *mut c_char;
557            
558            // Don't drop the vectors - Go will own the memory now
559            std::mem::forget(token_ids);
560            std::mem::forget(c_tokens);
561            
562            TokenizationResult {
563                token_ids: ids_ptr,
564                token_count: count,
565                tokens: tokens_ptr,
566                error: false
567            }
568        },
569        Err(e) => {
570            eprintln!("Error tokenizing text: {}", e);
571            TokenizationResult {
572                token_ids: std::ptr::null_mut(),
573                token_count: 0,
574                tokens: std::ptr::null_mut(),
575                error: true
576            }
577        }
578    }
579}
580
581// Free tokenization result allocated by Rust
582#[no_mangle]
583pub extern "C" fn free_tokenization_result(result: TokenizationResult) {
584    if !result.token_ids.is_null() && result.token_count > 0 {
585        unsafe {
586            // Reconstruct and drop the token_ids vector
587            let _ids_vec = Vec::from_raw_parts(result.token_ids, result.token_count as usize, result.token_count as usize);
588            
589            // Reconstruct and drop each token string
590            if !result.tokens.is_null() {
591                let tokens_slice = std::slice::from_raw_parts(result.tokens, result.token_count as usize);
592                for &token_ptr in tokens_slice {
593                    if !token_ptr.is_null() {
594                        let _ = CString::from_raw(token_ptr);
595                    }
596                }
597                
598                // Reconstruct and drop the tokens vector
599                let _tokens_vec = Vec::from_raw_parts(result.tokens, result.token_count as usize, result.token_count as usize);
600            }
601        }
602    }
603}
604
605// Initialize the BERT model (called from Go)
606#[no_mangle]
607pub extern "C" fn init_similarity_model(model_id: *const c_char, use_cpu: bool) -> bool {
608    let model_id = unsafe {
609        match CStr::from_ptr(model_id).to_str() {
610            Ok(s) => s,
611            Err(_) => return false,
612        }
613    };
614
615    match BertSimilarity::new(model_id, use_cpu) {
616        Ok(model) => {
617            let mut bert_opt = BERT_SIMILARITY.lock().unwrap();
618            *bert_opt = Some(model);
619            true
620        }
621        Err(e) => {
622            eprintln!("Failed to initialize BERT: {}", e);
623            false
624        }
625    }
626}
627
628// Structure to hold similarity result
629#[repr(C)]
630pub struct SimilarityResult {
631    pub index: i32,  // Index of the most similar text
632    pub score: f32,  // Similarity score
633}
634
635// Structure to hold embedding result
636#[repr(C)]
637pub struct EmbeddingResult {
638    pub data: *mut f32,
639    pub length: i32,
640    pub error: bool,
641}
642
643// Get embedding for a text (called from Go)
644#[no_mangle]
645pub extern "C" fn get_text_embedding(text: *const c_char, max_length: i32) -> EmbeddingResult {
646    let text = unsafe {
647        match CStr::from_ptr(text).to_str() {
648            Ok(s) => s,
649            Err(_) => return EmbeddingResult {
650                data: std::ptr::null_mut(),
651                length: 0,
652                error: true
653            },
654        }
655    };
656
657    let bert_opt = BERT_SIMILARITY.lock().unwrap();
658    let bert = match &*bert_opt {
659        Some(b) => b,
660        None => {
661            eprintln!("BERT model not initialized");
662            return EmbeddingResult {
663                data: std::ptr::null_mut(),
664                length: 0,
665                error: true
666            };
667        }
668    };
669
670    let max_length_opt = if max_length <= 0 { None } else { Some(max_length as usize) };
671    match bert.get_embedding(text, max_length_opt) {
672        Ok(embedding) => {
673            match embedding.flatten_all() {
674                Ok(flat_embedding) => {
675                    match flat_embedding.to_vec1::<f32>() {
676                        Ok(vec) => {
677                            let length = vec.len() as i32;
678                            // Allocate memory that will be freed by Go
679                            let data = vec.as_ptr() as *mut f32;
680                            std::mem::forget(vec); // Don't drop the vector - Go will own the memory now
681                            EmbeddingResult {
682                                data,
683                                length,
684                                error: false
685                            }
686                        },
687                        Err(_) => EmbeddingResult {
688                            data: std::ptr::null_mut(),
689                            length: 0,
690                            error: true
691                        }
692                    }
693                },
694                Err(_) => EmbeddingResult {
695                    data: std::ptr::null_mut(),
696                    length: 0,
697                    error: true
698                }
699            }
700        },
701        Err(e) => {
702            eprintln!("Error getting embedding: {}", e);
703            EmbeddingResult {
704                data: std::ptr::null_mut(),
705                length: 0,
706                error: true
707            }
708        }
709    }
710}
711
712// Calculate similarity between two texts (called from Go)
713#[no_mangle]
714pub extern "C" fn calculate_similarity(text1: *const c_char, text2: *const c_char, max_length: i32) -> f32 {
715    let text1 = unsafe {
716        match CStr::from_ptr(text1).to_str() {
717            Ok(s) => s,
718            Err(_) => return -1.0,
719        }
720    };
721    
722    let text2 = unsafe {
723        match CStr::from_ptr(text2).to_str() {
724            Ok(s) => s,
725            Err(_) => return -1.0,
726        }
727    };
728
729    let bert_opt = BERT_SIMILARITY.lock().unwrap();
730    let bert = match &*bert_opt {
731        Some(b) => b,
732        None => {
733            eprintln!("BERT model not initialized");
734            return -1.0;
735        }
736    };
737
738    let max_length_opt = if max_length <= 0 { None } else { Some(max_length as usize) };
739    match bert.calculate_similarity(text1, text2, max_length_opt) {
740        Ok(similarity) => similarity,
741        Err(e) => {
742            eprintln!("Error calculating similarity: {}", e);
743            -1.0
744        }
745    }
746}
747
748// Find most similar text from a list (called from Go)
749#[no_mangle]
750pub extern "C" fn find_most_similar(
751    query: *const c_char, 
752    candidates_ptr: *const *const c_char,
753    num_candidates: i32,
754    max_length: i32
755) -> SimilarityResult {
756    let query = unsafe {
757        match CStr::from_ptr(query).to_str() {
758            Ok(s) => s,
759            Err(_) => return SimilarityResult { index: -1, score: -1.0 },
760        }
761    };
762    
763    // Convert the array of C strings to Rust strings
764    let candidates: Vec<&str> = unsafe {
765        let mut result = Vec::with_capacity(num_candidates as usize);
766        let candidates_slice = std::slice::from_raw_parts(candidates_ptr, num_candidates as usize);
767        
768        for &cstr in candidates_slice {
769            match CStr::from_ptr(cstr).to_str() {
770                Ok(s) => result.push(s),
771                Err(_) => return SimilarityResult { index: -1, score: -1.0 },
772            }
773        }
774        
775        result
776    };
777
778    let bert_opt = BERT_SIMILARITY.lock().unwrap();
779    let bert = match &*bert_opt {
780        Some(b) => b,
781        None => {
782            eprintln!("BERT model not initialized");
783            return SimilarityResult { index: -1, score: -1.0 };
784        }
785    };
786
787    let max_length_opt = if max_length <= 0 { None } else { Some(max_length as usize) };
788    match bert.find_most_similar(query, &candidates, max_length_opt) {
789        Ok((idx, score)) => SimilarityResult { 
790            index: idx as i32, 
791            score 
792        },
793        Err(e) => {
794            eprintln!("Error finding most similar: {}", e);
795            SimilarityResult { index: -1, score: -1.0 }
796        }
797    }
798}
799
800// Free a C string allocated by Rust
801#[no_mangle]
802pub extern "C" fn free_cstring(s: *mut c_char) {
803    unsafe {
804        if !s.is_null() {
805            let _ = CString::from_raw(s);
806        }
807    }
808}
809
810// Free embedding data allocated by Rust
811#[no_mangle]
812pub extern "C" fn free_embedding(data: *mut f32, length: i32) {
813    if !data.is_null() && length > 0 {
814        unsafe {
815            // Reconstruct the vector so that Rust can properly deallocate it
816            let _vec = Vec::from_raw_parts(data, length as usize, length as usize);
817            // The vector will be dropped and the memory freed when _vec goes out of scope
818        }
819    }
820}
821
822// Helper function to L2 normalize a tensor
823fn normalize_l2(v: &Tensor) -> Result<Tensor> {
824    let norm = v.sqr()?.sum_keepdim(1)?.sqrt()?;
825    Ok(v.broadcast_div(&norm)?)
826}
827
828// New structure to hold classification result
829#[repr(C)]
830pub struct ClassificationResult {
831    pub class: i32,
832    pub confidence: f32,
833}
834
835// Initialize the BERT classifier model (called from Go)
836#[no_mangle]
837pub extern "C" fn init_classifier(model_id: *const c_char, num_classes: i32, use_cpu: bool) -> bool {
838    let model_id = unsafe {
839        match CStr::from_ptr(model_id).to_str() {
840            Ok(s) => s,
841            Err(_) => return false,
842        }
843    };
844
845    // Ensure num_classes is valid
846    if num_classes < 2 {
847        eprintln!("Number of classes must be at least 2, got {}", num_classes);
848        return false;
849    }
850
851    match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
852        Ok(classifier) => {
853            let mut bert_opt = BERT_CLASSIFIER.lock().unwrap();
854            *bert_opt = Some(classifier);
855            true
856        }
857        Err(e) => {
858            eprintln!("Failed to initialize BERT classifier: {}", e);
859            false
860        }
861    }
862}
863
864// Initialize the BERT PII classifier model (called from Go)
865#[no_mangle]
866pub extern "C" fn init_pii_classifier(model_id: *const c_char, num_classes: i32, use_cpu: bool) -> bool {
867    let model_id = unsafe {
868        match CStr::from_ptr(model_id).to_str() {
869            Ok(s) => s,
870            Err(_) => return false,
871        }
872    };
873
874    // Ensure num_classes is valid
875    if num_classes < 2 {
876        eprintln!("Number of classes must be at least 2, got {}", num_classes);
877        return false;
878    }
879
880    match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
881        Ok(classifier) => {
882            let mut bert_opt = BERT_PII_CLASSIFIER.lock().unwrap();
883            *bert_opt = Some(classifier);
884            true
885        }
886        Err(e) => {
887            eprintln!("Failed to initialize BERT PII classifier: {}", e);
888            false
889        }
890    }
891}
892
893// Initialize the BERT jailbreak classifier model (called from Go)
894#[no_mangle]
895pub extern "C" fn init_jailbreak_classifier(model_id: *const c_char, num_classes: i32, use_cpu: bool) -> bool {
896    let model_id = unsafe {
897        match CStr::from_ptr(model_id).to_str() {
898            Ok(s) => s,
899            Err(_) => return false,
900        }
901    };
902
903    // Ensure num_classes is valid
904    if num_classes < 2 {
905        eprintln!("Number of classes must be at least 2, got {}", num_classes);
906        return false;
907    }
908
909    match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
910        Ok(classifier) => {
911            let mut bert_opt = BERT_JAILBREAK_CLASSIFIER.lock().unwrap();
912            *bert_opt = Some(classifier);
913            true
914        }
915        Err(e) => {
916            eprintln!("Failed to initialize BERT jailbreak classifier: {}", e);
917            false
918        }
919    }
920}
921
922// Classify text using BERT (called from Go)
923#[no_mangle]
924pub extern "C" fn classify_text(text: *const c_char) -> ClassificationResult {
925    let default_result = ClassificationResult {
926        class: -1,
927        confidence: 0.0,
928    };
929
930    let text = unsafe {
931        match CStr::from_ptr(text).to_str() {
932            Ok(s) => s,
933            Err(_) => return default_result,
934        }
935    };
936
937    let bert_opt = BERT_CLASSIFIER.lock().unwrap();
938    match &*bert_opt {
939        Some(classifier) => match classifier.classify_text(text) {
940            Ok((class_idx, confidence)) => ClassificationResult {
941                class: class_idx as i32,
942                confidence,
943            },
944            Err(e) => {
945                eprintln!("Error classifying text: {}", e);
946                default_result
947            }
948        },
949        None => {
950            eprintln!("BERT classifier not initialized");
951            default_result
952        }
953    }
954}
955
956// Classify text for PII using BERT (called from Go)
957#[no_mangle]
958pub extern "C" fn classify_pii_text(text: *const c_char) -> ClassificationResult {
959    let default_result = ClassificationResult {
960        class: -1,
961        confidence: 0.0,
962    };
963
964    let text = unsafe {
965        match CStr::from_ptr(text).to_str() {
966            Ok(s) => s,
967            Err(_) => return default_result,
968        }
969    };
970
971    let bert_opt = BERT_PII_CLASSIFIER.lock().unwrap();
972    match &*bert_opt {
973        Some(classifier) => match classifier.classify_text(text) {
974            Ok((class_idx, confidence)) => ClassificationResult {
975                class: class_idx as i32,
976                confidence,
977            },
978            Err(e) => {
979                eprintln!("Error classifying PII text: {}", e);
980                default_result
981            }
982        },
983        None => {
984            eprintln!("BERT PII classifier not initialized");
985            default_result
986        }
987    }
988}
989
990// Classify text for jailbreak detection using BERT (called from Go)
991#[no_mangle]
992pub extern "C" fn classify_jailbreak_text(text: *const c_char) -> ClassificationResult {
993    let default_result = ClassificationResult {
994        class: -1,
995        confidence: 0.0,
996    };
997
998    let text = unsafe {
999        match CStr::from_ptr(text).to_str() {
1000            Ok(s) => s,
1001            Err(_) => return default_result,
1002        }
1003    };
1004
1005    let bert_opt = BERT_JAILBREAK_CLASSIFIER.lock().unwrap();
1006    match &*bert_opt {
1007        Some(classifier) => match classifier.classify_text(text) {
1008            Ok((class_idx, confidence)) => ClassificationResult {
1009                class: class_idx as i32,
1010                confidence,
1011            },
1012            Err(e) => {
1013                eprintln!("Error classifying jailbreak text: {}", e);
1014                default_result
1015            }
1016        },
1017        None => {
1018            eprintln!("BERT jailbreak classifier not initialized");
1019            default_result
1020        }
1021    }
1022}
candle_semantic_router/lib.rs

candle_semantic_router/
lib.rs
