1use std::collections::HashMap;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
14pub struct EntityMention {
15 pub text: String,
17 pub start: usize,
19 pub end: usize,
21}
22
23impl EntityMention {
24 pub fn new(text: impl Into<String>, start: usize, end: usize) -> Self {
26 Self {
27 text: text.into(),
28 start,
29 end,
30 }
31 }
32}
33
34#[derive(Debug, Clone)]
36pub struct EntityCandidate {
37 pub iri: String,
39 pub label: String,
41 pub score: f64,
43 pub aliases: Vec<String>,
45}
46
47impl EntityCandidate {
48 fn new(iri: impl Into<String>, label: impl Into<String>, aliases: Vec<String>) -> Self {
49 Self {
50 iri: iri.into(),
51 label: label.into(),
52 score: 0.0,
53 aliases,
54 }
55 }
56}
57
58#[derive(Debug, Clone)]
60pub struct LinkedEntity {
61 pub mention: EntityMention,
63 pub entity: EntityCandidate,
65 pub confidence: f64,
67}
68
69pub struct TfIdfIndex {
75 docs: Vec<(String, HashMap<String, f64>)>,
77 idf: HashMap<String, f64>,
79}
80
81impl TfIdfIndex {
82 pub fn new() -> Self {
84 Self {
85 docs: Vec::new(),
86 idf: HashMap::new(),
87 }
88 }
89
90 pub fn add_document(&mut self, doc_id: impl Into<String>, text: &str) {
92 let tokens = tokenize(text);
93 let total = tokens.len() as f64;
94 if total == 0.0 {
95 return;
96 }
97 let mut tf: HashMap<String, f64> = HashMap::new();
98 for tok in &tokens {
99 *tf.entry(tok.clone()).or_insert(0.0) += 1.0 / total;
100 }
101 self.docs.push((doc_id.into(), tf));
102 }
103
104 pub fn build(&mut self) {
106 let n = self.docs.len() as f64;
107 let mut df: HashMap<String, usize> = HashMap::new();
108 for (_, tf) in &self.docs {
109 for term in tf.keys() {
110 *df.entry(term.clone()).or_insert(0) += 1;
111 }
112 }
113 self.idf.clear();
114 for (term, count) in df {
115 self.idf.insert(term, (n / count as f64).ln() + 1.0);
116 }
117 }
118
119 pub fn similarity(&self, query: &str, doc_id: &str) -> f64 {
121 let doc = match self.docs.iter().find(|(id, _)| id == doc_id) {
122 Some((_, tf)) => tf,
123 None => return 0.0,
124 };
125
126 let q_tokens = tokenize(query);
127 let q_total = q_tokens.len() as f64;
128 if q_total == 0.0 {
129 return 0.0;
130 }
131 let mut q_tf: HashMap<String, f64> = HashMap::new();
132 for tok in &q_tokens {
133 *q_tf.entry(tok.clone()).or_insert(0.0) += 1.0 / q_total;
134 }
135
136 let mut dot = 0.0_f64;
138 let mut q_norm = 0.0_f64;
139 let mut d_norm = 0.0_f64;
140
141 let all_terms: std::collections::HashSet<&String> = q_tf.keys().chain(doc.keys()).collect();
142
143 for term in all_terms {
144 let idf = self.idf.get(term).copied().unwrap_or(1.0);
145 let q_val = q_tf.get(term).copied().unwrap_or(0.0) * idf;
146 let d_val = doc.get(term).copied().unwrap_or(0.0) * idf;
147 dot += q_val * d_val;
148 q_norm += q_val * q_val;
149 d_norm += d_val * d_val;
150 }
151
152 let denom = q_norm.sqrt() * d_norm.sqrt();
153 if denom < 1e-15 {
154 0.0
155 } else {
156 (dot / denom).clamp(0.0, 1.0)
157 }
158 }
159
160 pub fn doc_count(&self) -> usize {
162 self.docs.len()
163 }
164}
165
166impl Default for TfIdfIndex {
167 fn default() -> Self {
168 Self::new()
169 }
170}
171
172pub struct EntityLinker {
178 kb: HashMap<String, KbEntry>,
180 tfidf: TfIdfIndex,
182 first_char_index: HashMap<char, Vec<String>>,
188 pub nil_threshold: f64,
190}
191
192struct KbEntry {
193 label: String,
194 aliases: Vec<String>,
195}
196
197impl EntityLinker {
198 pub fn new() -> Self {
200 Self {
201 kb: HashMap::new(),
202 tfidf: TfIdfIndex::new(),
203 first_char_index: HashMap::new(),
204 nil_threshold: 0.1,
205 }
206 }
207
208 pub fn add_entity(
213 &mut self,
214 iri: impl Into<String>,
215 label: impl Into<String>,
216 aliases: &[&str],
217 ) {
218 let iri = iri.into();
219 let label = label.into();
220 let aliases: Vec<String> = aliases.iter().map(|s| s.to_string()).collect();
221 let context = format!("{} {}", label, aliases.join(" "));
222 self.tfidf.add_document(iri.clone(), &context);
223
224 let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
225 if let Some(c) = label.to_lowercase().chars().next() {
226 first_chars.insert(c);
227 }
228 for alias in &aliases {
229 if let Some(c) = alias.to_lowercase().chars().next() {
230 first_chars.insert(c);
231 }
232 }
233 for c in first_chars {
234 self.first_char_index
235 .entry(c)
236 .or_default()
237 .push(iri.clone());
238 }
239
240 self.kb.insert(iri, KbEntry { label, aliases });
241 }
242
243 pub fn build_index(&mut self) {
245 self.tfidf.build();
246 }
247
248 pub fn link(&self, text: &str) -> Vec<LinkedEntity> {
250 let mentions = detect_mentions(text);
251 let mut linked = Vec::new();
252
253 for mention in mentions {
254 let candidates = self.candidate_generation(&mention.text);
255 if candidates.is_empty() {
256 continue;
257 }
258 let best = self.disambiguate(&mention, &candidates, text);
259 if let Some(entity) = best {
260 let confidence = entity.score;
261 if confidence >= self.nil_threshold {
262 linked.push(LinkedEntity {
263 mention,
264 entity,
265 confidence,
266 });
267 }
268 }
269 }
270 linked
271 }
272
273 pub fn candidate_generation(&self, mention: &str) -> Vec<EntityCandidate> {
282 let mention_lower = mention.to_lowercase();
283 let Some(first_char) = mention_lower.chars().next() else {
284 return Vec::new();
285 };
286 let Some(blocked_iris) = self.first_char_index.get(&first_char) else {
287 return Vec::new();
288 };
289
290 let mut seen: std::collections::HashSet<&str> = std::collections::HashSet::new();
291 let mut candidates: Vec<EntityCandidate> = blocked_iris
292 .iter()
293 .filter(|iri| seen.insert(iri.as_str()))
294 .filter_map(|iri| {
295 let entry = self.kb.get(iri)?;
296 let label_score = jaro_winkler(&mention_lower, &entry.label.to_lowercase());
297 let alias_score = entry
298 .aliases
299 .iter()
300 .map(|a| jaro_winkler(&mention_lower, &a.to_lowercase()))
301 .fold(0.0_f64, f64::max);
302 let score = label_score.max(alias_score);
303 if score > 0.6 {
304 let mut c = EntityCandidate::new(
305 iri.clone(),
306 entry.label.clone(),
307 entry.aliases.clone(),
308 );
309 c.score = score;
310 Some(c)
311 } else {
312 None
313 }
314 })
315 .collect();
316
317 candidates.sort_by(|a, b| {
318 b.score
319 .partial_cmp(&a.score)
320 .unwrap_or(std::cmp::Ordering::Equal)
321 });
322 candidates
323 }
324
325 pub fn disambiguate(
327 &self,
328 _mention: &EntityMention,
329 candidates: &[EntityCandidate],
330 context: &str,
331 ) -> Option<EntityCandidate> {
332 if candidates.is_empty() {
333 return None;
334 }
335
336 let mut best_score = f64::NEG_INFINITY;
337 let mut best: Option<EntityCandidate> = None;
338
339 for cand in candidates {
340 let ctx_score = self.tfidf.similarity(context, &cand.iri);
341 let combined = cand.score * 0.6 + ctx_score * 0.4;
343 if combined > best_score {
344 best_score = combined;
345 let mut winner = cand.clone();
346 winner.score = combined;
347 best = Some(winner);
348 }
349 }
350 best
351 }
352
353 pub fn entity_count(&self) -> usize {
355 self.kb.len()
356 }
357}
358
359impl Default for EntityLinker {
360 fn default() -> Self {
361 Self::new()
362 }
363}
364
365fn detect_mentions(text: &str) -> Vec<EntityMention> {
372 let mut mentions = Vec::new();
373 let mut chars = text.char_indices().peekable();
374 let bytes = text.as_bytes();
375 let len = bytes.len();
376
377 while let Some((start, ch)) = chars.next() {
378 if ch.is_uppercase() {
379 let mut end = start + ch.len_utf8();
381 while end < len {
382 let next_ch = text[end..].chars().next().unwrap_or('\0');
383 if next_ch.is_alphanumeric() || next_ch == ' ' {
384 if next_ch == ' ' {
386 let after_space = end + 1;
387 if after_space < len {
388 let nc2 = text[after_space..].chars().next().unwrap_or('\0');
389 if nc2.is_uppercase() {
390 end = after_space + nc2.len_utf8();
391 let _ = chars.next(); let _ = chars.next(); continue;
395 }
396 }
397 break;
398 }
399 end += next_ch.len_utf8();
400 let _ = chars.next();
401 } else {
402 break;
403 }
404 }
405 let mention_text = text[start..end].trim().to_string();
406 if mention_text.len() >= 2 {
407 mentions.push(EntityMention::new(mention_text, start, end));
408 }
409 }
410 }
411 mentions
412}
413
414fn jaro_winkler(s1: &str, s2: &str) -> f64 {
416 if s1 == s2 {
417 return 1.0;
418 }
419 let jaro = jaro(s1, s2);
420 let prefix_len = s1
421 .chars()
422 .zip(s2.chars())
423 .take(4)
424 .take_while(|(a, b)| a == b)
425 .count();
426 let p = 0.1_f64;
427 jaro + (prefix_len as f64 * p * (1.0 - jaro))
428}
429
430fn jaro(s1: &str, s2: &str) -> f64 {
431 let s1: Vec<char> = s1.chars().collect();
432 let s2: Vec<char> = s2.chars().collect();
433 let len1 = s1.len();
434 let len2 = s2.len();
435 if len1 == 0 && len2 == 0 {
436 return 1.0;
437 }
438 if len1 == 0 || len2 == 0 {
439 return 0.0;
440 }
441
442 let match_window = (len1.max(len2) / 2).saturating_sub(1);
443 let mut s1_matches = vec![false; len1];
444 let mut s2_matches = vec![false; len2];
445 let mut matches = 0usize;
446 let mut transpositions = 0usize;
447
448 for (i, &c1) in s1.iter().enumerate() {
449 let start = i.saturating_sub(match_window);
450 let end = (i + match_window + 1).min(len2);
451 for (j, &c2) in s2[start..end].iter().enumerate() {
452 let j_real = start + j;
453 if !s2_matches[j_real] && c1 == c2 {
454 s1_matches[i] = true;
455 s2_matches[j_real] = true;
456 matches += 1;
457 break;
458 }
459 }
460 }
461
462 if matches == 0 {
463 return 0.0;
464 }
465
466 let mut k = 0;
467 for (i, &s1m) in s1_matches.iter().enumerate() {
468 if s1m {
469 while !s2_matches[k] {
470 k += 1;
471 }
472 if s1[i] != s2[k] {
473 transpositions += 1;
474 }
475 k += 1;
476 }
477 }
478
479 let m = matches as f64;
480 (m / len1 as f64 + m / len2 as f64 + (m - transpositions as f64 / 2.0) / m) / 3.0
481}
482
483fn tokenize(text: &str) -> Vec<String> {
485 text.split(|c: char| !c.is_alphanumeric())
486 .filter(|s| !s.is_empty())
487 .map(|s| s.to_lowercase())
488 .collect()
489}
490
491#[cfg(test)]
496mod tests {
497 use super::*;
498
499 fn linker_with_persons() -> EntityLinker {
500 let mut linker = EntityLinker::new();
501 linker.add_entity(
502 "http://example.org/Albert_Einstein",
503 "Albert Einstein",
504 &["Einstein", "A. Einstein"],
505 );
506 linker.add_entity(
507 "http://example.org/Marie_Curie",
508 "Marie Curie",
509 &["Curie", "M. Curie"],
510 );
511 linker.add_entity(
512 "http://example.org/Isaac_Newton",
513 "Isaac Newton",
514 &["Newton"],
515 );
516 linker.build_index();
517 linker
518 }
519
520 #[test]
523 fn test_mention_new() {
524 let m = EntityMention::new("Alice", 0, 5);
525 assert_eq!(m.text, "Alice");
526 assert_eq!(m.start, 0);
527 assert_eq!(m.end, 5);
528 }
529
530 #[test]
531 fn test_mention_equality() {
532 let m1 = EntityMention::new("Bob", 0, 3);
533 let m2 = EntityMention::new("Bob", 0, 3);
534 assert_eq!(m1, m2);
535 }
536
537 #[test]
540 fn test_tfidf_add_document() {
541 let mut idx = TfIdfIndex::new();
542 idx.add_document("doc1", "quantum physics relativity");
543 idx.build();
544 assert_eq!(idx.doc_count(), 1);
545 }
546
547 #[test]
548 fn test_tfidf_similarity_same_doc() {
549 let mut idx = TfIdfIndex::new();
550 idx.add_document("doc1", "quantum physics relativity");
551 idx.build();
552 let sim = idx.similarity("quantum physics", "doc1");
553 assert!(sim > 0.0, "similarity should be > 0, got {sim}");
554 }
555
556 #[test]
557 fn test_tfidf_similarity_different_content() {
558 let mut idx = TfIdfIndex::new();
559 idx.add_document("doc1", "quantum physics relativity");
560 idx.add_document("doc2", "cooking recipes baking bread");
561 idx.build();
562 let s1 = idx.similarity("quantum physics", "doc1");
563 let s2 = idx.similarity("quantum physics", "doc2");
564 assert!(s1 > s2, "physics query should match doc1 better");
565 }
566
567 #[test]
568 fn test_tfidf_unknown_doc() {
569 let idx = TfIdfIndex::new();
570 assert_eq!(idx.similarity("anything", "unknown"), 0.0);
571 }
572
573 #[test]
574 fn test_tfidf_empty_query() {
575 let mut idx = TfIdfIndex::new();
576 idx.add_document("d", "hello world");
577 idx.build();
578 assert_eq!(idx.similarity("", "d"), 0.0);
579 }
580
581 #[test]
582 fn test_tfidf_default() {
583 let idx = TfIdfIndex::default();
584 assert_eq!(idx.doc_count(), 0);
585 }
586
587 #[test]
590 fn test_linker_entity_count() {
591 let linker = linker_with_persons();
592 assert_eq!(linker.entity_count(), 3);
593 }
594
595 #[test]
596 fn test_linker_default() {
597 let linker = EntityLinker::default();
598 assert_eq!(linker.entity_count(), 0);
599 }
600
601 #[test]
604 fn test_candidate_generation_exact_label() {
605 let linker = linker_with_persons();
606 let cands = linker.candidate_generation("Einstein");
607 assert!(!cands.is_empty());
608 assert!(cands[0].iri.contains("Einstein"));
609 }
610
611 #[test]
612 fn test_candidate_generation_partial() {
613 let linker = linker_with_persons();
614 let cands = linker.candidate_generation("Newton");
615 assert!(!cands.is_empty());
616 assert!(cands.iter().any(|c| c.iri.contains("Newton")));
617 }
618
619 #[test]
620 fn test_candidate_generation_no_match() {
621 let linker = linker_with_persons();
622 let cands = linker.candidate_generation("Zorkblat");
623 assert!(cands.is_empty());
624 }
625
626 #[test]
627 fn test_candidate_generation_sorted_by_score() {
628 let linker = linker_with_persons();
629 let cands = linker.candidate_generation("Curie");
630 for i in 1..cands.len() {
631 assert!(cands[i - 1].score >= cands[i].score);
632 }
633 }
634
635 #[test]
636 fn test_candidate_generation_alias_match() {
637 let linker = linker_with_persons();
638 let cands = linker.candidate_generation("Curie");
640 assert!(cands.iter().any(|c| c.iri.contains("Curie")));
641 }
642
643 #[test]
647 fn test_candidate_generation_blocking_scales_with_many_entities() {
648 let mut linker = EntityLinker::new();
649 for i in 0..300 {
650 linker.add_entity(
651 format!("http://example.org/Other{i}"),
652 format!("Other{i}"),
653 &[],
654 );
655 }
656 linker.add_entity(
657 "http://example.org/Albert_Einstein",
658 "Albert Einstein",
659 &["Einstein", "A. Einstein"],
660 );
661 linker.build_index();
662
663 let cands = linker.candidate_generation("Einstein");
664 assert!(!cands.is_empty());
665 assert!(cands[0].iri.contains("Einstein"));
666 }
667
668 #[test]
671 fn test_candidate_generation_blocking_no_match_for_unindexed_first_char() {
672 let linker = linker_with_persons();
673 let cands = linker.candidate_generation("Zzzzzzz");
674 assert!(cands.is_empty());
675 }
676
677 #[test]
680 fn test_disambiguate_returns_best() {
681 let linker = linker_with_persons();
682 let cands = linker.candidate_generation("Einstein");
683 let mention = EntityMention::new("Einstein", 0, 8);
684 let best = linker.disambiguate(&mention, &cands, "Einstein worked on relativity");
685 assert!(best.is_some());
686 assert!(best.expect("should succeed").iri.contains("Einstein"));
687 }
688
689 #[test]
690 fn test_disambiguate_empty_candidates() {
691 let linker = linker_with_persons();
692 let mention = EntityMention::new("X", 0, 1);
693 let result = linker.disambiguate(&mention, &[], "context");
694 assert!(result.is_none());
695 }
696
697 #[test]
698 fn test_disambiguate_score_in_range() {
699 let linker = linker_with_persons();
700 let cands = linker.candidate_generation("Newton");
701 let mention = EntityMention::new("Newton", 0, 6);
702 if let Some(best) = linker.disambiguate(&mention, &cands, "gravity laws Newton") {
703 assert!((0.0..=1.0).contains(&best.score));
704 }
705 }
706
707 #[test]
710 fn test_link_finds_entity() {
711 let linker = linker_with_persons();
712 let linked = linker.link("Einstein developed relativity theory.");
713 assert!(!linked.is_empty());
714 assert!(linked[0].entity.iri.contains("Einstein"));
715 }
716
717 #[test]
718 fn test_link_confidence_above_threshold() {
719 let linker = linker_with_persons();
720 let linked = linker.link("Newton formulated laws of motion.");
721 for le in &linked {
722 assert!(le.confidence >= linker.nil_threshold);
723 }
724 }
725
726 #[test]
727 fn test_link_no_entities_in_empty_text() {
728 let linker = linker_with_persons();
729 let linked = linker.link("");
730 assert!(linked.is_empty());
731 }
732
733 #[test]
734 fn test_link_result_fields() {
735 let linker = linker_with_persons();
736 let linked = linker.link("Einstein and Curie were scientists.");
737 for le in &linked {
738 assert!(!le.mention.text.is_empty());
739 assert!(!le.entity.iri.is_empty());
740 assert!((0.0..=1.0).contains(&le.confidence));
741 }
742 }
743
744 #[test]
747 fn test_jaro_winkler_identical() {
748 assert!((jaro_winkler("hello", "hello") - 1.0).abs() < 1e-9);
749 }
750
751 #[test]
752 fn test_jaro_winkler_completely_different() {
753 let score = jaro_winkler("abc", "xyz");
754 assert!(score < 0.5, "score = {score}");
755 }
756
757 #[test]
758 fn test_jaro_winkler_prefix_boost() {
759 let jw = jaro_winkler("einstein", "einstien");
760 assert!(jw > 0.8, "score = {jw}");
761 }
762
763 #[test]
764 fn test_jaro_winkler_empty_strings() {
765 assert!((jaro("", "") - 1.0).abs() < 1e-9);
766 assert!((jaro("abc", "") - 0.0).abs() < 1e-9);
767 }
768
769 #[test]
772 fn test_detect_mentions_finds_capitalized() {
773 let mentions = detect_mentions("Alice and Bob went to Paris.");
774 let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
775 assert!(texts
777 .iter()
778 .any(|t| *t == "Alice" || t.starts_with("Alice")));
779 }
780
781 #[test]
782 fn test_detect_mentions_empty() {
783 assert!(detect_mentions("").is_empty());
784 }
785
786 #[test]
787 fn test_detect_mentions_lowercase_only() {
788 let mentions = detect_mentions("all lowercase words here");
789 assert!(mentions.is_empty());
790 }
791
792 #[test]
795 fn test_tokenize_basic() {
796 let tokens = tokenize("Hello World");
797 assert_eq!(tokens, vec!["hello", "world"]);
798 }
799
800 #[test]
801 fn test_tokenize_empty() {
802 assert!(tokenize("").is_empty());
803 }
804
805 #[test]
806 fn test_tokenize_punctuation_split() {
807 let tokens = tokenize("foo, bar; baz.");
808 assert_eq!(tokens, vec!["foo", "bar", "baz"]);
809 }
810
811 #[test]
814 fn test_full_pipeline() {
815 let mut linker = EntityLinker::new();
816 linker.add_entity("http://ex.org/Paris", "Paris", &["City of Light"]);
817 linker.add_entity("http://ex.org/London", "London", &["British capital"]);
818 linker.build_index();
819
820 let linked = linker.link("Paris is a famous city in France.");
821 if !linked.is_empty() {
822 assert!(linked[0].entity.iri.contains("Paris"));
823 }
824 }
826
827 #[test]
828 fn test_nil_threshold_filters_low_confidence() {
829 let mut linker = EntityLinker::new();
830 linker.add_entity("http://ex.org/X", "Xyzzy", &[]);
831 linker.build_index();
832 linker.nil_threshold = 0.99; let linked = linker.link("Xyzzy something");
835 for le in &linked {
837 assert!(le.confidence >= 0.99);
838 }
839 }
840}