Skip to main content

gaze_recognizers/ner/
recognizer.rs

1use std::path::Path;
2
3use gaze_types::{Candidate, ConflictTier, DetectContext, PiiClass, Recognizer};
4
5#[cfg(feature = "test-support")]
6use super::backend::test_support::load_test_support_recognizer;
7use super::detector::NerDetector;
8use super::error::NerLoadError;
9use super::types::NerOptions;
10
11pub struct NerRecognizer {
12    pub(crate) detector: NerDetector,
13}
14
15impl NerRecognizer {
16    pub fn load_with_options(model_dir: &Path, options: NerOptions) -> Result<Self, NerLoadError> {
17        #[cfg(feature = "test-support")]
18        if let Some(recognizer) = load_test_support_recognizer(model_dir, &options) {
19            return Ok(recognizer);
20        }
21
22        Ok(Self {
23            detector: NerDetector::load_with_options(model_dir, options)?,
24        })
25    }
26}
27
28impl Recognizer for NerRecognizer {
29    fn id(&self) -> &str {
30        "ner"
31    }
32
33    fn supported_class(&self) -> &PiiClass {
34        &PiiClass::Name
35    }
36
37    fn detect(&self, input: &str, _ctx: &DetectContext<'_>) -> Vec<Candidate> {
38        match self.detector.backend.detect(input) {
39            Ok(spans) => spans
40                .into_iter()
41                .filter(|span| span.score >= self.detector.threshold)
42                .map(|span| {
43                    Candidate::new(
44                        span.span,
45                        span.class,
46                        self.id(),
47                        span.score,
48                        0,
49                        None,
50                        self.token_family(),
51                        format!("ner/{}", self.detector.backend_kind.as_str()),
52                        ConflictTier::None,
53                        Vec::new(),
54                    )
55                    .with_recognizer_version_id(self.detector.recognizer_version_id())
56                })
57                .collect(),
58            Err(err) => {
59                tracing::warn!(backend = self.detector.backend_kind.as_str(), error = %err, "ner: backend detect failed");
60                Vec::new()
61            }
62        }
63    }
64
65    fn token_family(&self) -> &str {
66        "counter"
67    }
68}