Skip to content

Commit 01edddf

Browse files
rootfsrootfs
committed
update tests
Signed-off-by: Huamin Chen <[email protected]>
1 parent 9b36254 commit 01edddf

File tree

2 files changed

+201
-631
lines changed

2 files changed

+201
-631
lines changed

‎candle-binding/src/ffi/init.rs‎

Lines changed: 17 additions & 85 deletions
Original file line numberDiff line numberDiff line change
@@ -136,14 +136,8 @@ pub extern "C" fn init_similarity_model(model_id: *const c_char, use_cpu: bool)
136136

137137
match BertSimilarity::new(model_id, use_cpu) {
138138
Ok(model) => {
139-
// Set using OnceLock - fails if already initialized
140-
match BERT_SIMILARITY.set(Arc::new(model)) {
141-
Ok(_) => true,
142-
Err(_) => {
143-
eprintln!("BERT similarity model already initialized");
144-
false
145-
}
146-
}
139+
// Set using OnceLock - returns false if already initialized (safe to re-call)
140+
BERT_SIMILARITY.set(Arc::new(model)).is_ok()
147141
}
148142
Err(e) => {
149143
eprintln!("Failed to initialize BERT: {e}");
@@ -177,13 +171,7 @@ pub extern "C" fn init_classifier(
177171
}
178172

179173
match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
180-
Ok(classifier) => match BERT_CLASSIFIER.set(Arc::new(classifier)) {
181-
Ok(_) => true,
182-
Err(_) => {
183-
eprintln!("BERT classifier already initialized");
184-
false
185-
}
186-
},
174+
Ok(classifier) => BERT_CLASSIFIER.set(Arc::new(classifier)).is_ok(),
187175
Err(e) => {
188176
eprintln!("Failed to initialize BERT classifier: {e}");
189177
false
@@ -215,13 +203,7 @@ pub extern "C" fn init_pii_classifier(
215203
}
216204

217205
match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
218-
Ok(classifier) => match BERT_PII_CLASSIFIER.set(Arc::new(classifier)) {
219-
Ok(_) => true,
220-
Err(_) => {
221-
eprintln!("BERT PII classifier already initialized");
222-
false
223-
}
224-
},
206+
Ok(classifier) => BERT_PII_CLASSIFIER.set(Arc::new(classifier)).is_ok(),
225207
Err(e) => {
226208
eprintln!("Failed to initialize BERT PII classifier: {e}");
227209
false
@@ -253,13 +235,7 @@ pub extern "C" fn init_jailbreak_classifier(
253235
}
254236

255237
match BertClassifier::new(model_id, num_classes as usize, use_cpu) {
256-
Ok(classifier) => match BERT_JAILBREAK_CLASSIFIER.set(Arc::new(classifier)) {
257-
Ok(_) => true,
258-
Err(_) => {
259-
eprintln!("BERT jailbreak classifier already initialized");
260-
false
261-
}
262-
},
238+
Ok(classifier) => BERT_JAILBREAK_CLASSIFIER.set(Arc::new(classifier)).is_ok(),
263239
Err(e) => {
264240
eprintln!("Failed to initialize BERT jailbreak classifier: {e}");
265241
false
@@ -283,13 +259,9 @@ pub extern "C" fn init_modernbert_classifier(model_id: *const c_char, use_cpu: b
283259
// Try to initialize the actual ModernBERT model using traditional architecture
284260
match crate::model_architectures::traditional::modernbert::TraditionalModernBertClassifier::load_from_directory(model_id, use_cpu) {
285261
Ok(model) => {
286-
match crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_CLASSIFIER.set(Arc::new(model)) {
287-
Ok(_) => true,
288-
Err(_) => {
289-
eprintln!("ModernBERT classifier already initialized");
290-
false
291-
}
292-
}
262+
crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_CLASSIFIER
263+
.set(Arc::new(model))
264+
.is_ok()
293265
}
294266
Err(e) => {
295267
eprintln!("Failed to initialize ModernBERT classifier: {}", e);
@@ -314,13 +286,7 @@ pub extern "C" fn init_modernbert_pii_classifier(model_id: *const c_char, use_cp
314286
// Try to initialize the actual ModernBERT PII model
315287
match crate::model_architectures::traditional::modernbert::TraditionalModernBertClassifier::load_from_directory(model_id, use_cpu) {
316288
Ok(model) => {
317-
match crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_PII_CLASSIFIER.set(Arc::new(model)) {
318-
Ok(_) => true,
319-
Err(_) => {
320-
eprintln!("ModernBERT PII classifier already initialized");
321-
false
322-
}
323-
}
289+
crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_PII_CLASSIFIER.set(Arc::new(model)).is_ok()
324290
}
325291
Err(e) => {
326292
eprintln!("Failed to initialize ModernBERT PII classifier: {}", e);
@@ -349,13 +315,7 @@ pub extern "C" fn init_modernbert_pii_token_classifier(
349315
// Create the token classifier
350316
match crate::model_architectures::traditional::modernbert::TraditionalModernBertTokenClassifier::new(model_id, use_cpu) {
351317
Ok(classifier) => {
352-
match crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_TOKEN_CLASSIFIER.set(Arc::new(classifier)) {
353-
Ok(_) => true,
354-
Err(_) => {
355-
eprintln!("ModernBERT PII token classifier already initialized");
356-
false
357-
}
358-
}
318+
crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_TOKEN_CLASSIFIER.set(Arc::new(classifier)).is_ok()
359319
}
360320
Err(e) => {
361321
println!(" ERROR: Failed to initialize ModernBERT PII token classifier: {}", e);
@@ -383,13 +343,7 @@ pub extern "C" fn init_modernbert_jailbreak_classifier(
383343
// Try to initialize the actual ModernBERT jailbreak model
384344
match crate::model_architectures::traditional::modernbert::TraditionalModernBertClassifier::load_from_directory(model_id, use_cpu) {
385345
Ok(model) => {
386-
match crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_JAILBREAK_CLASSIFIER.set(Arc::new(model)) {
387-
Ok(_) => true,
388-
Err(_) => {
389-
eprintln!("ModernBERT jailbreak classifier already initialized");
390-
false
391-
}
392-
}
346+
crate::model_architectures::traditional::modernbert::TRADITIONAL_MODERNBERT_JAILBREAK_CLASSIFIER.set(Arc::new(model)).is_ok()
393347
}
394348
Err(e) => {
395349
eprintln!("Failed to initialize ModernBERT jailbreak classifier: {}", e);
@@ -506,13 +460,7 @@ pub extern "C" fn init_unified_classifier_c(
506460
Ok(mut classifier) => {
507461
// Initialize traditional path with actual models
508462
match classifier.init_traditional_path() {
509-
Ok(_) => match UNIFIED_CLASSIFIER.set(Arc::new(classifier)) {
510-
Ok(_) => true,
511-
Err(_) => {
512-
eprintln!("Unified classifier already initialized");
513-
false
514-
}
515-
},
463+
Ok(_) => UNIFIED_CLASSIFIER.set(Arc::new(classifier)).is_ok(),
516464
Err(e) => {
517465
eprintln!("Failed to initialize traditional path: {}", e);
518466
false
@@ -638,13 +586,7 @@ pub extern "C" fn init_candle_bert_token_classifier(
638586
match crate::classifiers::lora::token_lora::LoRATokenClassifier::new(
639587
model_path, use_cpu,
640588
) {
641-
Ok(classifier) => match LORA_TOKEN_CLASSIFIER.set(Arc::new(classifier)) {
642-
Ok(_) => true,
643-
Err(_) => {
644-
eprintln!("LoRA token classifier already initialized");
645-
false
646-
}
647-
},
589+
Ok(classifier) => LORA_TOKEN_CLASSIFIER.set(Arc::new(classifier)).is_ok(),
648590
Err(e) => {
649591
eprintln!(" ERROR: Failed to initialize LoRA token classifier: {}", e);
650592
false
@@ -659,13 +601,9 @@ pub extern "C" fn init_candle_bert_token_classifier(
659601
use_cpu,
660602
) {
661603
Ok(classifier) => {
662-
match crate::model_architectures::traditional::bert::TRADITIONAL_BERT_TOKEN_CLASSIFIER.set(Arc::new(classifier)) {
663-
Ok(_) => true,
664-
Err(_) => {
665-
eprintln!("Traditional BERT token classifier already initialized");
666-
false
667-
}
668-
}
604+
crate::model_architectures::traditional::bert::TRADITIONAL_BERT_TOKEN_CLASSIFIER
605+
.set(Arc::new(classifier))
606+
.is_ok()
669607
}
670608
Err(e) => {
671609
eprintln!(
@@ -760,13 +698,7 @@ pub extern "C" fn init_lora_unified_classifier(
760698
) {
761699
Ok(engine) => {
762700
// Store in global static variable (Arc for efficient cloning during concurrent access)
763-
match PARALLEL_LORA_ENGINE.set(Arc::new(engine)) {
764-
Ok(_) => true,
765-
Err(_) => {
766-
eprintln!("LoRA parallel engine already initialized");
767-
false
768-
}
769-
}
701+
PARALLEL_LORA_ENGINE.set(Arc::new(engine)).is_ok()
770702
}
771703
Err(e) => {
772704
eprintln!(

0 commit comments

Comments
 (0)