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