gaze_recognizers/ner/
recognizer.rs1use 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}