Skip to main content

anno_eval/eval/
backend_factory.rs

1//! Backend Factory for Runtime Backend Creation
2//!
3//! This module provides a factory pattern for creating backend instances
4//! from string names, enabling dynamic backend selection for evaluation.
5//!
6//! # Design Philosophy
7//!
8//! - **Feature-aware**: Only creates backends when features are enabled
9//! - **Graceful degradation**: Returns errors for unavailable backends
10//! - **Model defaults**: Uses sensible default models for each backend
11//! - **Trait-based**: Returns trait objects for polymorphic usage
12
13use anno::{Model, Result};
14
15/// Factory for creating backend instances from names.
16pub struct BackendFactory;
17
18impl BackendFactory {
19    /// Create a backend instance from a name.
20    ///
21    /// # Supported Backends
22    ///
23    /// ## Always Available
24    /// - `pattern` / `RegexNER` - Pattern-based NER
25    /// - `heuristic` / `HeuristicNER` - Heuristic NER
26    /// - `stacked` / `StackedNER` - Stacked NER
27    ///
28    /// ## ONNX Feature Required
29    /// - `bert_onnx` / `BertNEROnnx` - BERT ONNX NER
30    /// - `gliner_onnx` / `GLiNEROnnx` - GLiNER ONNX (zero-shot)
31    /// - `nuner` / `NuNER` - NuNER (zero-shot, token-based)
32    /// - `w2ner` / `W2NER` - W2NER (discontinuous NER)
33    /// - `gliner_multitask` / `GLiNERMultitaskOnnx` - GLiNERMultitask multi-task
34    ///
35    /// ## Candle Feature Required
36    /// - `candle_ner` / `CandleNER` - Candle BERT NER
37    /// - `gliner_candle` / `GLiNERCandle` - GLiNER Candle (zero-shot)
38    /// - `gliner_multitask_candle` / `GLiNERMultitaskCandle` - GLiNERMultitask Candle
39    ///
40    /// ## Coreference
41    /// - `coref_resolver` / `SimpleCorefResolver` - Simple coreference resolver
42    ///
43    /// # Example
44    ///
45    /// ```rust,ignore
46    /// use anno_eval::eval::backend_factory::BackendFactory;
47    ///
48    /// let backend = BackendFactory::create("pattern")?;
49    /// let entities = backend.extract_entities("Meeting on Jan 15", None)?;
50    /// ```
51    pub fn create(backend_name: &str) -> Result<Box<dyn Model>> {
52        match backend_name.to_lowercase().as_str() {
53            // Always available backends
54            "pattern" | "patternner" | "regex" | "regexner" => Ok(Box::new(anno::RegexNER::new())),
55            "heuristic" | "heuristicner" => Ok(Box::new(anno::HeuristicNER::new())),
56            "stacked" | "stackedner" => Ok(Box::new(anno::StackedNER::default())),
57            "crf" | "crfner" => Ok(Box::new(anno::backends::crf::CrfNER::new())),
58            "hmm" | "hmmner" => Ok(Box::new(anno::backends::hmm::HmmNER::new())),
59            "ensemble" | "ensemblener" => {
60                use anno::backends::ensemble::EnsembleNER;
61                Ok(Box::new(EnsembleNER::default()) as Box<dyn Model>)
62            }
63            "heuristic_crf" | "heuristic-crf" | "heuristiccrfner" => {
64                use anno::backends::heuristic_crf::HeuristicCrfNER;
65                Ok(Box::new(HeuristicCrfNER::new()) as Box<dyn Model>)
66            }
67            #[cfg(feature = "heuristic-fr")]
68            "heuristic_fr" | "heuristic-fr" | "heuristicfrner" => {
69                use anno::backends::heuristic_fr::HeuristicFrNer;
70                Ok(Box::new(HeuristicFrNer::new()) as Box<dyn Model>)
71            }
72
73            // ONNX backends
74            #[cfg(feature = "onnx")]
75            "bert_onnx" | "bertneronnx" => {
76                use anno::backends::onnx::BertNEROnnx;
77                use crate::DEFAULT_BERT_ONNX_MODEL;
78                BertNEROnnx::new(DEFAULT_BERT_ONNX_MODEL)
79                    .map(|m| Box::new(m) as Box<dyn Model>)
80                    .map_err(|e| {
81                        anno::Error::FeatureNotAvailable(format!(
82                            "Failed to create BertNEROnnx: {}",
83                            e
84                        ))
85                    })
86            }
87            #[cfg(not(feature = "onnx"))]
88            "bert_onnx" | "bertneronnx" => Err(anno::Error::FeatureNotAvailable(
89                "BertNEROnnx requires 'onnx' feature".to_string(),
90            )),
91
92            #[cfg(feature = "onnx")]
93            "gliner" => {
94                // First-class alias: prefer ONNX when available.
95                use anno::backends::gliner_onnx::GLiNEROnnx;
96                use crate::DEFAULT_GLINER_MODEL;
97                GLiNEROnnx::new(DEFAULT_GLINER_MODEL)
98                    .map(|m| Box::new(m) as Box<dyn Model>)
99                    .map_err(|e| {
100                        anno::Error::FeatureNotAvailable(format!(
101                            "Failed to create GLiNER (onnx): {}",
102                            e
103                        ))
104                    })
105            }
106            #[cfg(all(not(feature = "onnx"), feature = "candle"))]
107            "gliner" => {
108                // Fallback alias: Candle implementation when ONNX isn't enabled.
109                use anno::backends::gliner_candle::GLiNERCandle;
110                use crate::DEFAULT_GLINER_CANDLE_MODEL;
111                GLiNERCandle::from_pretrained(DEFAULT_GLINER_CANDLE_MODEL)
112                    .map(|m| Box::new(m) as Box<dyn Model>)
113                    .map_err(|e| {
114                        anno::Error::FeatureNotAvailable(format!(
115                            "Failed to create GLiNER (candle): {}",
116                            e
117                        ))
118                    })
119            }
120            #[cfg(all(not(feature = "onnx"), not(feature = "candle")))]
121            "gliner" => Err(crate::Error::FeatureNotAvailable(
122                "GLiNER requires 'onnx' (preferred) or 'candle' feature".to_string(),
123            )),
124
125            #[cfg(feature = "onnx")]
126            "gliner_onnx" | "glineronnx" => {
127                use crate::backends::gliner_onnx::GLiNEROnnx;
128                use crate::DEFAULT_GLINER_MODEL;
129                GLiNEROnnx::new(DEFAULT_GLINER_MODEL)
130                    .map(|m| Box::new(m) as Box<dyn Model>)
131                    .map_err(|e| {
132                        crate::Error::FeatureNotAvailable(format!(
133                            "Failed to create GLiNEROnnx: {}",
134                            e
135                        ))
136                    })
137            }
138            #[cfg(not(feature = "onnx"))]
139            "gliner_onnx" | "glineronnx" => Err(crate::Error::FeatureNotAvailable(
140                "GLiNEROnnx requires 'onnx' feature".to_string(),
141            )),
142
143            // B2NER (COLING 2025, unified NER training on 54 datasets)
144            // Note: only LLM-scale models (7B/20B) available on HuggingFace.
145            // Requires LLM backend, not ONNX. Pending encoder-scale release.
146            "b2ner" => Err(crate::Error::FeatureNotAvailable(
147                "B2NER only has LLM-scale models (7B/20B) on HuggingFace. \
148                 Encoder-scale ONNX weights pending release."
149                    .to_string(),
150            )),
151
152            // GLiNER PII Edge (60+ PII categories, zero-shot)
153            #[cfg(feature = "onnx")]
154            "gliner_pii" | "pii_ml" => {
155                use crate::backends::gliner_onnx::GLiNEROnnx;
156                GLiNEROnnx::new(anno::models::GLINER_PII)
157                    .map(|m| Box::new(m) as Box<dyn Model>)
158                    .map_err(|e| {
159                        crate::Error::FeatureNotAvailable(format!(
160                            "GLiNER PII Edge model unavailable: {}",
161                            e
162                        ))
163                    })
164            }
165            #[cfg(not(feature = "onnx"))]
166            "gliner_pii" | "pii_ml" => Err(crate::Error::FeatureNotAvailable(
167                "GLiNER PII requires 'onnx' feature".to_string(),
168            )),
169
170            // GLiNER-RelEx (joint NER + relation extraction, zero-shot)
171            #[cfg(feature = "onnx")]
172            "gliner_relex" | "relex" => {
173                use crate::backends::gliner_onnx::GLiNEROnnx;
174                GLiNEROnnx::new(anno::models::GLINER_RELEX)
175                    .map(|m| Box::new(m) as Box<dyn Model>)
176                    .map_err(|e| {
177                        crate::Error::FeatureNotAvailable(format!(
178                            "GLiNER-RelEx model unavailable: {}\n\n\
179                             Export: uv run scripts/export_glirel_onnx.py",
180                            e
181                        ))
182                    })
183            }
184            #[cfg(not(feature = "onnx"))]
185            "gliner_relex" | "relex" => Err(crate::Error::FeatureNotAvailable(
186                "GLiNER-RelEx requires 'onnx' feature".to_string(),
187            )),
188
189            #[cfg(feature = "onnx")]
190            "nuner" | "nunerzero" => {
191                use crate::backends::nuner::NuNER;
192                use crate::DEFAULT_NUNER_MODEL;
193                NuNER::from_pretrained(DEFAULT_NUNER_MODEL)
194                    .map(|m| Box::new(m) as Box<dyn Model>)
195                    .map_err(|e| {
196                        crate::Error::FeatureNotAvailable(format!("Failed to create NuNER: {}", e))
197                    })
198            }
199            #[cfg(not(feature = "onnx"))]
200            "nuner" | "nunerzero" => Err(crate::Error::FeatureNotAvailable(
201                "NuNER requires 'onnx' feature".to_string(),
202            )),
203
204            #[cfg(feature = "onnx")]
205            "nuner_4k" | "nunerzero4k" => {
206                use crate::backends::nuner::NuNER;
207                NuNER::from_pretrained("numind/NuNER_Zero-4k")
208                    .map(|m| Box::new(m) as Box<dyn Model>)
209                    .map_err(|e| {
210                        crate::Error::FeatureNotAvailable(format!(
211                            "Failed to create NuNER 4k: {}",
212                            e
213                        ))
214                    })
215            }
216            #[cfg(not(feature = "onnx"))]
217            "nuner_4k" | "nunerzero4k" => Err(crate::Error::FeatureNotAvailable(
218                "NuNER 4k requires 'onnx' feature".to_string(),
219            )),
220
221            #[cfg(feature = "onnx")]
222            "w2ner" => {
223                use crate::backends::w2ner::W2NER;
224                use crate::DEFAULT_W2NER_MODEL;
225                // Allow override via environment variable for custom/exported models
226                let model_path = std::env::var("W2NER_MODEL_PATH")
227                    .unwrap_or_else(|_| DEFAULT_W2NER_MODEL.to_string());
228                W2NER::from_pretrained(&model_path)
229                    .map(|m| Box::new(m) as Box<dyn Model>)
230                    .map_err(|e| {
231                        crate::Error::FeatureNotAvailable(format!(
232                            "W2NER model unavailable: {}\n\n\
233                             Options:\n\
234                             1. Set W2NER_MODEL_PATH to a local model directory\n\
235                             2. Export your own: uv run scripts/export_w2ner_to_onnx.py\n\
236                             3. For HF models, set HF_TOKEN and request access",
237                            e
238                        ))
239                    })
240            }
241            #[cfg(not(feature = "onnx"))]
242            "w2ner" => Err(crate::Error::FeatureNotAvailable(
243                "W2NER requires 'onnx' feature".to_string(),
244            )),
245
246            #[cfg(feature = "onnx")]
247            "gliner_multitask" | "gliner_multitask_onnx" => {
248                use crate::backends::gliner_multitask::GLiNERMultitaskOnnx;
249                use crate::DEFAULT_GLINER_MULTITASK_MODEL;
250                GLiNERMultitaskOnnx::from_pretrained(DEFAULT_GLINER_MULTITASK_MODEL)
251                    .map(|m| Box::new(m) as Box<dyn Model>)
252                    .map_err(|e| {
253                        crate::Error::FeatureNotAvailable(format!(
254                            "Failed to create GLiNER multi-task (ONNX): {}",
255                            e
256                        ))
257                    })
258            }
259            #[cfg(not(feature = "onnx"))]
260            "gliner_multitask" | "gliner_multitask_onnx" => Err(crate::Error::FeatureNotAvailable(
261                "GLiNER multi-task (ONNX) requires 'onnx' feature".to_string(),
262            )),
263
264            // Candle backends
265            #[cfg(feature = "candle")]
266            "candle_ner" | "candlener" => {
267                use crate::backends::candle::CandleNER;
268                use crate::DEFAULT_CANDLE_MODEL;
269                CandleNER::from_pretrained(DEFAULT_CANDLE_MODEL)
270                    .map(|m| Box::new(m) as Box<dyn Model>)
271                    .map_err(|e| {
272                        crate::Error::FeatureNotAvailable(format!(
273                            "CandleNER model unavailable: {}",
274                            e
275                        ))
276                    })
277            }
278            #[cfg(not(feature = "candle"))]
279            "candle_ner" | "candlener" => Err(crate::Error::FeatureNotAvailable(
280                "CandleNER requires 'candle' feature".to_string(),
281            )),
282
283            #[cfg(feature = "candle")]
284            "gliner_candle" | "glinercandle" => {
285                use crate::backends::gliner_candle::GLiNERCandle;
286                use crate::DEFAULT_GLINER_CANDLE_MODEL;
287                GLiNERCandle::from_pretrained(DEFAULT_GLINER_CANDLE_MODEL)
288                    .map(|m| Box::new(m) as Box<dyn Model>)
289                    .map_err(|e| {
290                        crate::Error::FeatureNotAvailable(format!(
291                            "GLiNERCandle model unavailable: {}",
292                            e
293                        ))
294                    })
295            }
296            #[cfg(not(feature = "candle"))]
297            "gliner_candle" | "glinercandle" => Err(crate::Error::FeatureNotAvailable(
298                "GLiNERCandle requires 'candle' feature".to_string(),
299            )),
300
301            #[cfg(all(feature = "candle", feature = "onnx"))]
302            "gliner_multitask_candle" => {
303                use crate::backends::gliner_multitask::GLiNERMultitaskCandle;
304                use crate::DEFAULT_GLINER_MULTITASK_MODEL;
305                GLiNERMultitaskCandle::from_pretrained(DEFAULT_GLINER_MULTITASK_MODEL)
306                    .map(|m| Box::new(m) as Box<dyn Model>)
307                    .map_err(|e| {
308                        crate::Error::FeatureNotAvailable(format!(
309                            "Failed to create GLiNER multi-task (Candle): {}",
310                            e
311                        ))
312                    })
313            }
314            #[cfg(not(all(feature = "candle", feature = "onnx")))]
315            "gliner_multitask_candle" => Err(crate::Error::FeatureNotAvailable(
316                "GLiNER multi-task (Candle) requires both 'candle' and 'onnx' features".to_string(),
317            )),
318
319            // TPLinker (ONNX neural with `onnx` feature, heuristic fallback otherwise)
320            "tplinker" | "tplink" => {
321                use anno::backends::tplinker::TPLinker;
322                Ok(Box::new(TPLinker::new()?) as Box<dyn Model>)
323            }
324
325            // Poly-Encoder GLiNER (requires onnx)
326            #[cfg(feature = "onnx")]
327            "gliner_poly" | "gliner-poly" | "poly_gliner" => {
328                use anno::backends::gliner_poly::GLiNERPoly;
329                use anno::DEFAULT_GLINER_POLY_MODEL;
330                GLiNERPoly::new(DEFAULT_GLINER_POLY_MODEL)
331                    .map(|m| Box::new(m) as Box<dyn anno::Model>)
332                    .map_err(|e| crate::Error::model_init(e.to_string()))
333            }
334            #[cfg(not(feature = "onnx"))]
335            "gliner_poly" | "gliner-poly" | "poly_gliner" => Err(crate::Error::FeatureNotAvailable(
336                "GLiNERPoly requires 'onnx' feature".to_string(),
337            )),
338
339            // DeBERTa-v3 NER (requires onnx) -- uses BertNEROnnx (same ONNX interface)
340            #[cfg(feature = "onnx")]
341            "deberta_v3" | "deberta-v3" | "deberta" => {
342                use anno::backends::onnx::BertNEROnnx;
343                let Ok(model_path) = std::env::var("DEBERTA_MODEL_PATH") else {
344                    return Err(crate::Error::FeatureNotAvailable(
345                        "DeBERTa-v3 backend requires a local ONNX export. Set DEBERTA_MODEL_PATH (e.g. after running `uv run scripts/export_deberta_ner_to_onnx.py`)."
346                            .to_string(),
347                    ));
348                };
349                BertNEROnnx::new(&model_path)
350                    .map(|m| Box::new(m) as Box<dyn Model>)
351                    .map_err(|e| {
352                        crate::Error::Retrieval(format!(
353                            "DeBERTa-v3 model unavailable: {e}\n\n\
354                             Options:\n\
355                             1. Export your own: uv run scripts/export_deberta_ner_to_onnx.py\n\
356                             2. Set DEBERTA_MODEL_PATH to a local model directory\n\
357                             3. Use --model bert-onnx or --model candle-ner instead",
358                        ))
359                    })
360            }
361            #[cfg(not(feature = "onnx"))]
362            "deberta_v3" | "deberta-v3" | "deberta" => Err(crate::Error::FeatureNotAvailable(
363                "DeBERTa-v3 NER requires 'onnx' feature".to_string(),
364            )),
365
366            // ALBERT NER (requires onnx) -- uses BertNEROnnx (same ONNX interface)
367            #[cfg(feature = "onnx")]
368            "albert" | "albert_ner" => {
369                use anno::backends::onnx::BertNEROnnx;
370                let Ok(model_path) = std::env::var("ALBERT_MODEL_PATH") else {
371                    return Err(crate::Error::FeatureNotAvailable(
372                        "ALBERT backend requires a local ONNX export. Set ALBERT_MODEL_PATH to a local model directory containing ONNX weights."
373                            .to_string(),
374                    ));
375                };
376                BertNEROnnx::new(&model_path)
377                    .map(|m| Box::new(m) as Box<dyn Model>)
378                    .map_err(|e| {
379                        crate::Error::Retrieval(format!(
380                            "ALBERT model unavailable: {e}\n\n\
381                             Options:\n\
382                             1. Export your own ONNX model\n\
383                             2. Set ALBERT_MODEL_PATH to a local model directory\n\
384                             3. Use --model bert-onnx or --model candle-ner instead",
385                        ))
386                    })
387            }
388            #[cfg(not(feature = "onnx"))]
389            "albert" | "albert_ner" => Err(crate::Error::FeatureNotAvailable(
390                "ALBERT NER requires 'onnx' feature".to_string(),
391            )),
392
393            // UniversalNER (LLM-backed zero-shot, requires `llm` feature + API key)
394            "universal_ner" | "universal-ner" | "universalner" => {
395                use anno::backends::universal_ner::UniversalNER;
396                let m = UniversalNER::new()?;
397                if !m.is_available() {
398                    return Err(crate::Error::FeatureNotAvailable(
399                        "UniversalNER requires the `llm` feature and a non-empty API key. Set one of: OPENAI_API_KEY, ANTHROPIC_API_KEY, OPENROUTER_API_KEY, GEMINI_API_KEY, or UNIVERSAL_NER_API_KEY."
400                            .to_string(),
401                    ));
402                }
403                Ok(Box::new(m) as Box<dyn Model>)
404            }
405
406            // Unknown backend
407            _ => Err(crate::Error::InvalidInput(format!(
408                "Unknown backend: '{}'. Available: pattern, heuristic, stacked, crf, hmm, ensemble, heuristic_crf, tplinker{}",
409                backend_name,
410                if cfg!(feature = "onnx") {
411                    ", bert_onnx, gliner_onnx, nuner, w2ner, gliner_multitask"
412                } else {
413                    ""
414                }
415            ))),
416        }
417    }
418
419    /// List all available backends (based on enabled features).
420    #[must_use]
421    pub fn available_backends() -> Vec<&'static str> {
422        #[allow(unused_mut)] // mut needed for extend/push calls
423        let mut backends = vec![
424            "pattern",
425            "heuristic",
426            "stacked",
427            "crf",
428            "hmm",
429            "ensemble",
430            "heuristic_crf",
431            "tplinker",
432        ];
433
434        // UniversalNER requires the optional `llm` feature plus a non-empty API key.
435        // If either is missing, treat it as unavailable to avoid “Feature not available”
436        // failures in the matrix harness.
437        if cfg!(feature = "llm") {
438            anno::env::load_dotenv();
439            if anno::env::has_llm_api_key() || std::env::var("UNIVERSAL_NER_API_KEY").is_ok() {
440                backends.push("universal_ner");
441            }
442        }
443
444        #[cfg(feature = "onnx")]
445        {
446            backends.extend(&[
447                "bert_onnx",
448                "gliner",
449                "gliner_onnx",
450                "nuner",
451                "nuner_4k",
452                "b2ner",
453                "w2ner",
454                "gliner_multitask",
455                "gliner_pii",
456                "gliner_relex",
457                "gliner_poly",
458            ]);
459
460            // Optional backends that require explicit local ONNX exports.
461            if std::env::var("DEBERTA_MODEL_PATH").is_ok() {
462                backends.push("deberta_v3");
463            }
464            if std::env::var("ALBERT_MODEL_PATH").is_ok() {
465                backends.push("albert");
466            }
467        }
468
469        #[cfg(feature = "candle")]
470        {
471            backends.extend(&["candle_ner", "gliner_candle"]);
472            // `gliner` is also available as an alias when candle is enabled
473            // (and onnx is not required).
474            if !cfg!(feature = "onnx") {
475                backends.push("gliner");
476            }
477        }
478
479        #[cfg(all(feature = "candle", feature = "onnx"))]
480        {
481            backends.push("gliner_multitask_candle");
482        }
483
484        backends
485    }
486
487    /// List all available coreference resolvers.
488    ///
489    /// Coreference resolvers are *not* `Model`s, so they are kept separate from
490    /// [`Self::available_backends`]. They are used by `TaskEvaluator` for coref-family tasks.
491    #[must_use]
492    pub fn available_coref_resolvers() -> Vec<&'static str> {
493        vec!["coref_resolver", "mention_ranking"]
494    }
495
496    /// Check if a backend is available (feature-enabled).
497    #[must_use]
498    pub fn is_available(backend_name: &str) -> bool {
499        Self::available_backends().contains(&backend_name.to_lowercase().as_str())
500    }
501}
502
503/// Helper to create a coreference resolver from a name.
504///
505/// Note: Coreference resolvers don't implement `Model`, so this is separate.
506pub fn create_coref_resolver(
507    name: &str,
508) -> Result<Box<dyn crate::eval::coref_resolver::CoreferenceResolver>> {
509    match name.to_lowercase().as_str() {
510        "coref_resolver" | "simplecorefresolver" | "simple" => {
511            use crate::eval::coref_resolver::{CorefConfig, SimpleCorefResolver};
512            Ok(Box::new(SimpleCorefResolver::new(CorefConfig::default())))
513        }
514        "mention_ranking" | "mention-ranking" | "mentionranking" => {
515            use anno::backends::coref::mention_ranking::MentionRankingCoref;
516            Ok(Box::new(MentionRankingCoref::new()))
517        }
518        _ => Err(crate::Error::InvalidInput(format!(
519            "Unknown coreference resolver: '{}'. Available: coref_resolver, mention_ranking",
520            name
521        ))),
522    }
523}
524
525/// Create a text-based coreference backend (CorefBackend trait).
526///
527/// Unlike `create_coref_resolver` which takes pre-extracted entities,
528/// `CorefBackend` operates on raw text and returns mention clusters directly.
529/// This is the interface used by neural coref models (FCoref, MentionRanking).
530pub fn create_coref_backend(name: &str) -> Result<Box<dyn anno::CorefBackend>> {
531    match name.to_lowercase().as_str() {
532        "mention_ranking" | "mention-ranking" | "mentionranking" => {
533            use anno::backends::coref::mention_ranking::MentionRankingCoref;
534            Ok(Box::new(MentionRankingCoref::new()))
535        }
536        #[cfg(feature = "onnx")]
537        "fcoref" | "f-coref" | "fastcoref" => {
538            use anno::backends::coref::fcoref::FCoref;
539            let model_path = std::env::var("FCOREF_MODEL_PATH").ok();
540            let fcoref = if let Some(path) = model_path {
541                FCoref::from_path(&path)?
542            } else {
543                FCoref::from_pretrained("biu-nlp/f-coref")?
544            };
545            Ok(Box::new(fcoref))
546        }
547        #[cfg(not(feature = "onnx"))]
548        "fcoref" | "f-coref" | "fastcoref" => Err(crate::Error::FeatureNotAvailable(
549            "FCoref requires 'onnx' feature. Export: uv run scripts/export_fcoref.py".to_string(),
550        )),
551        _ => Err(crate::Error::InvalidInput(format!(
552            "Unknown coref backend: '{}'. Available: mention_ranking, fcoref",
553            name
554        ))),
555    }
556}
557
558/// List available coref backends (text-based CorefBackend).
559pub fn available_coref_backends() -> Vec<&'static str> {
560    #[allow(unused_mut)]
561    let mut backends = vec!["mention_ranking"];
562    #[cfg(feature = "onnx")]
563    {
564        backends.push("fcoref");
565    }
566    backends
567}
568
569#[cfg(test)]
570mod tests {
571    use super::*;
572
573    #[test]
574    fn known_backends_construct_with_expected_name_prefix() {
575        // Stacked reports a composite name like "stacked(regex+heuristic)", so
576        // match against a prefix per backend rather than exact-equality.
577        let cases = [
578            ("pattern", "regex"),
579            ("heuristic", "heuristic"),
580            ("stacked", "stacked"),
581        ];
582        for (alias, expected_prefix) in cases {
583            let model = BackendFactory::create(alias).unwrap();
584            let name = model.name();
585            assert!(
586                name.starts_with(expected_prefix),
587                "for alias {alias:?}: expected prefix {expected_prefix:?}, got {name:?}"
588            );
589        }
590    }
591
592    #[test]
593    fn test_unknown_backend() {
594        let backend = BackendFactory::create("nonexistent");
595        assert!(backend.is_err());
596    }
597
598    #[test]
599    fn test_available_backends() {
600        let backends = BackendFactory::available_backends();
601        assert!(backends.contains(&"pattern"));
602        assert!(backends.contains(&"heuristic"));
603        assert!(backends.contains(&"stacked"));
604    }
605}
606
607#[cfg(test)]
608mod additional_tests {
609    use super::*;
610
611    #[test]
612    fn test_backend_factory_pattern_returns_regex_only() {
613        let model = BackendFactory::create("pattern").unwrap();
614        println!("Model name: {}", model.name());
615        assert_eq!(model.name(), "regex", "pattern should return RegexNER");
616
617        let entities = model
618            .extract_entities("John Smith went to Paris", None)
619            .unwrap();
620        println!("Entities: {:?}", entities);
621
622        // Should NOT have PER or LOC
623        for e in &entities {
624            assert!(
625                !matches!(e.entity_type, crate::EntityType::Person),
626                "Unexpected Person entity: {:?}",
627                e
628            );
629            assert!(
630                !matches!(e.entity_type, crate::EntityType::Location),
631                "Unexpected Location entity: {:?}",
632                e
633            );
634        }
635    }
636}