gaze_recognizers/ner/
recognizer.rs

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

gaze_recognizers/ner/
recognizer.rs
