Skip to main content

gaze_recognizers/ner/
recognizer.rs

1use std::path::Path;
2
3use gaze_types::{Candidate, ConflictTier, DetectContext, DetectError, 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<'_>) -> Result<Vec<Candidate>, DetectError> {
38        self.detector
39            .detect_span_results(input)
40            .map_err(|err| {
41                tracing::error!(backend = self.detector.backend_kind.as_str(), error = %err, "ner: backend detect failed closed");
42                DetectError::backend(self.id(), err.to_string())
43            })
44            .map(|spans| {
45                spans
46                    .into_iter()
47                    .filter(|span| span.score >= self.detector.threshold)
48                    .map(|span| {
49                        Candidate::new(
50                            span.span,
51                            span.class,
52                            self.id(),
53                            span.score,
54                            0,
55                            None,
56                            self.token_family(),
57                            format!("ner/{}", self.detector.backend_kind.as_str()),
58                            ConflictTier::None,
59                            Vec::new(),
60                        )
61                        .with_recognizer_version_id(self.detector.recognizer_version_id())
62                    })
63                    .collect()
64            })
65    }
66
67    fn token_family(&self) -> &str {
68        "counter"
69    }
70}