1use std::collections::HashMap;
7
8#[derive(Debug, Clone)]
14pub struct Entity {
15 pub iri: String,
17 pub label: String,
19 pub aliases: Vec<String>,
21 pub entity_type: String,
23 pub description: Option<String>,
25 pub popularity: f64,
27}
28
29#[derive(Debug, Clone)]
31pub struct EntityMention {
32 pub text: String,
34 pub start_char: usize,
36 pub end_char: usize,
38 pub candidates: Vec<LinkCandidate>,
40}
41
42#[derive(Debug, Clone)]
44pub struct LinkCandidate {
45 pub entity: Entity,
47 pub score: f64,
49 pub string_similarity: f64,
51 pub prior_probability: f64,
53}
54
55#[derive(Debug, Clone)]
57pub struct LinkedEntity {
58 pub mention: EntityMention,
60 pub best_candidate: Option<LinkCandidate>,
62 pub confidence: f64,
64}
65
66#[derive(Debug)]
72pub enum LinkerError {
73 EmptyText,
75 EmptyKnowledgeBase,
77 InvalidEntity(String),
79}
80
81impl std::fmt::Display for LinkerError {
82 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
83 match self {
84 LinkerError::EmptyText => write!(f, "Input text is empty"),
85 LinkerError::EmptyKnowledgeBase => write!(f, "Knowledge base contains no entities"),
86 LinkerError::InvalidEntity(msg) => write!(f, "Invalid entity: {}", msg),
87 }
88 }
89}
90
91impl std::error::Error for LinkerError {}
92
93pub struct EntityLinker {
99 entities: Vec<Entity>,
100 label_index: HashMap<String, Vec<usize>>,
102 first_char_index: HashMap<char, Vec<usize>>,
107 max_label_chars: usize,
110 pub min_mention_len: usize,
112 pub threshold: f64,
114}
115
116impl EntityLinker {
117 pub fn new(threshold: f64) -> Self {
123 Self {
124 entities: Vec::new(),
125 label_index: HashMap::new(),
126 first_char_index: HashMap::new(),
127 max_label_chars: 0,
128 min_mention_len: 2,
129 threshold,
130 }
131 }
132
133 pub fn add_entity(&mut self, entity: Entity) {
135 let idx = self.entities.len();
136 let mut keys: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
138 keys.push(entity.label.to_lowercase());
139
140 let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
141 for key in &keys {
142 self.max_label_chars = self.max_label_chars.max(key.chars().count());
143 if let Some(c) = key.chars().next() {
144 first_chars.insert(c);
145 }
146 }
147 for c in first_chars {
148 self.first_char_index.entry(c).or_default().push(idx);
149 }
150
151 for key in keys {
152 self.label_index.entry(key).or_default().push(idx);
153 }
154 self.entities.push(entity);
155 }
156
157 pub fn add_entities(&mut self, entities: Vec<Entity>) {
159 for entity in entities {
160 self.add_entity(entity);
161 }
162 }
163
164 pub fn entity_count(&self) -> usize {
166 self.entities.len()
167 }
168
169 pub fn detect_mentions(&self, text: &str) -> Vec<EntityMention> {
184 let text_lower = text.to_lowercase();
185 let mut mentions: Vec<EntityMention> = Vec::new();
186
187 if self.label_index.is_empty() || self.max_label_chars == 0 {
188 return mentions;
189 }
190
191 let boundaries: Vec<usize> = text_lower
195 .char_indices()
196 .map(|(byte_idx, _)| byte_idx)
197 .chain(std::iter::once(text_lower.len()))
198 .collect();
199 let num_chars = boundaries.len() - 1;
200
201 for start in 0..num_chars {
202 let start_byte = boundaries[start];
203 let max_len = self.max_label_chars.min(num_chars - start);
204
205 for len in self.min_mention_len..=max_len {
206 let end_char = start + len;
207 let end_byte = boundaries[end_char];
208 let surface = &text_lower[start_byte..end_byte];
209
210 let Some(indices) = self.label_index.get(surface) else {
211 continue;
212 };
213
214 let candidates = self.candidates_for_surface(surface, indices);
215 if !candidates.is_empty() {
216 let original_text = &text[start_byte..end_byte];
218 mentions.push(EntityMention {
219 text: original_text.to_string(),
220 start_char: start_byte,
221 end_char: end_byte,
222 candidates,
223 });
224 }
225 }
226 }
227
228 mentions.sort_by_key(|m| (m.start_char, usize::MAX - (m.end_char - m.start_char)));
230 let mut deduped: Vec<EntityMention> = Vec::new();
231 for mention in mentions {
232 if let Some(last) = deduped.last() {
234 if mention.start_char >= last.start_char && mention.end_char <= last.end_char {
235 continue;
236 }
237 }
238 deduped.push(mention);
239 }
240 deduped
241 }
242
243 pub fn link(&self, text: &str) -> Result<Vec<LinkedEntity>, LinkerError> {
253 if text.trim().is_empty() {
254 return Err(LinkerError::EmptyText);
255 }
256 if self.entities.is_empty() {
257 return Err(LinkerError::EmptyKnowledgeBase);
258 }
259
260 let mentions = self.detect_mentions(text);
261 let linked = mentions
262 .into_iter()
263 .map(|mention| {
264 let best = mention
265 .candidates
266 .iter()
267 .find(|c| c.score >= self.threshold)
268 .cloned();
269 let confidence = best.as_ref().map(|c| c.score).unwrap_or(0.0);
270 LinkedEntity {
271 mention,
272 best_candidate: best,
273 confidence,
274 }
275 })
276 .collect();
277 Ok(linked)
278 }
279
280 pub fn link_mention(&self, mention: &str) -> Vec<LinkCandidate> {
288 let mention_lower = mention.to_lowercase();
289 let Some(first_char) = mention_lower.chars().next() else {
290 return Vec::new();
291 };
292 let Some(blocked_indices) = self.first_char_index.get(&first_char) else {
293 return Vec::new();
294 };
295
296 let mut seen: std::collections::HashSet<usize> = std::collections::HashSet::new();
297 let mut candidates: Vec<LinkCandidate> = Vec::new();
298
299 for &idx in blocked_indices {
300 if !seen.insert(idx) {
301 continue;
302 }
303 let Some(entity) = self.entities.get(idx) else {
304 continue;
305 };
306
307 let mut surfaces: Vec<String> =
309 entity.aliases.iter().map(|a| a.to_lowercase()).collect();
310 surfaces.push(entity.label.to_lowercase());
311
312 let best_sim = surfaces
313 .iter()
314 .map(|s| Self::string_similarity(&mention_lower, s))
315 .fold(0.0_f64, f64::max);
316
317 if best_sim > 0.0 {
318 let score = Self::combined_score(best_sim, entity.popularity);
319 candidates.push(LinkCandidate {
320 entity: entity.clone(),
321 score,
322 string_similarity: best_sim,
323 prior_probability: entity.popularity,
324 });
325 }
326 }
327
328 candidates.sort_by(|a, b| {
330 b.score
331 .partial_cmp(&a.score)
332 .unwrap_or(std::cmp::Ordering::Equal)
333 });
334 candidates
335 }
336
337 pub fn find_by_iri(&self, iri: &str) -> Option<&Entity> {
339 self.entities.iter().find(|e| e.iri == iri)
340 }
341
342 pub fn string_similarity(a: &str, b: &str) -> f64 {
351 if a == b {
352 return 1.0;
353 }
354 let len_a = a.chars().count();
355 let len_b = b.chars().count();
356 let max_len = len_a.max(len_b);
357 if max_len == 0 {
358 return 1.0; }
360 let dist = Self::edit_distance(a, b);
361 1.0 - (dist as f64 / max_len as f64)
362 }
363
364 fn edit_distance(a: &str, b: &str) -> usize {
366 let a_chars: Vec<char> = a.chars().collect();
367 let b_chars: Vec<char> = b.chars().collect();
368 let la = a_chars.len();
369 let lb = b_chars.len();
370
371 if la == 0 {
372 return lb;
373 }
374 if lb == 0 {
375 return la;
376 }
377
378 let mut dp = vec![vec![0usize; lb + 1]; la + 1];
379 for (i, row) in dp.iter_mut().enumerate() {
380 row[0] = i;
381 }
382 for (j, cell) in dp[0].iter_mut().enumerate() {
383 *cell = j;
384 }
385 for i in 1..=la {
386 for j in 1..=lb {
387 let cost = if a_chars[i - 1] == b_chars[j - 1] {
388 0
389 } else {
390 1
391 };
392 dp[i][j] = (dp[i - 1][j] + 1)
393 .min(dp[i][j - 1] + 1)
394 .min(dp[i - 1][j - 1] + cost);
395 }
396 }
397 dp[la][lb]
398 }
399
400 #[allow(dead_code)]
406 fn rebuild_index(&mut self) {
407 self.label_index.clear();
408 self.first_char_index.clear();
409 self.max_label_chars = 0;
410 for (idx, entity) in self.entities.iter().enumerate() {
411 let mut keys: Vec<String> = entity.aliases.iter().map(|a| a.to_lowercase()).collect();
412 keys.push(entity.label.to_lowercase());
413
414 let mut first_chars: std::collections::HashSet<char> = std::collections::HashSet::new();
415 for key in &keys {
416 self.max_label_chars = self.max_label_chars.max(key.chars().count());
417 if let Some(c) = key.chars().next() {
418 first_chars.insert(c);
419 }
420 }
421 for c in first_chars {
422 self.first_char_index.entry(c).or_default().push(idx);
423 }
424
425 for key in keys {
426 self.label_index.entry(key).or_default().push(idx);
427 }
428 }
429 }
430
431 fn candidates_for_surface(&self, surface: &str, indices: &[usize]) -> Vec<LinkCandidate> {
433 let mut candidates: Vec<LinkCandidate> = indices
434 .iter()
435 .filter_map(|&idx| {
436 let entity = self.entities.get(idx)?;
437 let sim = Self::string_similarity(surface, &entity.label.to_lowercase());
438 let score = Self::combined_score(sim, entity.popularity);
439 Some(LinkCandidate {
440 entity: entity.clone(),
441 score,
442 string_similarity: sim,
443 prior_probability: entity.popularity,
444 })
445 })
446 .collect();
447 candidates.sort_by(|a, b| {
448 b.score
449 .partial_cmp(&a.score)
450 .unwrap_or(std::cmp::Ordering::Equal)
451 });
452 candidates
453 }
454
455 fn combined_score(sim: f64, popularity: f64) -> f64 {
457 0.7 * sim + 0.3 * popularity
458 }
459}
460
461#[cfg(test)]
466mod tests {
467 use super::*;
468
469 fn make_entity(iri: &str, label: &str, aliases: &[&str], pop: f64) -> Entity {
470 Entity {
471 iri: iri.to_string(),
472 label: label.to_string(),
473 aliases: aliases.iter().map(|s| s.to_string()).collect(),
474 entity_type: "Thing".to_string(),
475 description: None,
476 popularity: pop,
477 }
478 }
479
480 fn sample_linker() -> EntityLinker {
481 let mut linker = EntityLinker::new(0.3);
482 linker.add_entity(make_entity(
483 "http://example.org/einstein",
484 "Albert Einstein",
485 &["Einstein", "A. Einstein"],
486 0.95,
487 ));
488 linker.add_entity(make_entity(
489 "http://example.org/curie",
490 "Marie Curie",
491 &["Curie", "M. Curie"],
492 0.90,
493 ));
494 linker.add_entity(make_entity(
495 "http://example.org/berlin",
496 "Berlin",
497 &["Berlin City"],
498 0.80,
499 ));
500 linker
501 }
502
503 #[test]
506 fn test_entity_count_empty() {
507 let linker = EntityLinker::new(0.5);
508 assert_eq!(linker.entity_count(), 0);
509 }
510
511 #[test]
512 fn test_entity_count_after_add() {
513 let mut linker = EntityLinker::new(0.5);
514 linker.add_entity(make_entity("http://x.org/a", "Alpha", &[], 0.5));
515 assert_eq!(linker.entity_count(), 1);
516 linker.add_entity(make_entity("http://x.org/b", "Beta", &[], 0.5));
517 assert_eq!(linker.entity_count(), 2);
518 }
519
520 #[test]
521 fn test_add_entities_bulk() {
522 let mut linker = EntityLinker::new(0.5);
523 linker.add_entities(vec![
524 make_entity("http://x.org/a", "Alpha", &[], 0.5),
525 make_entity("http://x.org/b", "Beta", &[], 0.6),
526 make_entity("http://x.org/c", "Gamma", &[], 0.7),
527 ]);
528 assert_eq!(linker.entity_count(), 3);
529 }
530
531 #[test]
534 fn test_find_by_iri_exists() {
535 let linker = sample_linker();
536 let entity = linker.find_by_iri("http://example.org/einstein");
537 assert!(entity.is_some());
538 assert_eq!(entity.expect("some").label, "Albert Einstein");
539 }
540
541 #[test]
542 fn test_find_by_iri_not_found() {
543 let linker = sample_linker();
544 assert!(linker.find_by_iri("http://example.org/nobody").is_none());
545 }
546
547 #[test]
550 fn test_string_similarity_exact() {
551 assert!((EntityLinker::string_similarity("hello", "hello") - 1.0).abs() < 1e-9);
552 }
553
554 #[test]
555 fn test_string_similarity_completely_different() {
556 let sim = EntityLinker::string_similarity("abc", "xyz");
557 assert!(sim < 1.0);
558 }
559
560 #[test]
561 fn test_string_similarity_both_empty() {
562 assert!((EntityLinker::string_similarity("", "") - 1.0).abs() < 1e-9);
563 }
564
565 #[test]
566 fn test_string_similarity_one_empty() {
567 let sim = EntityLinker::string_similarity("", "hello");
568 assert!((sim - 0.0).abs() < 1e-9);
569 }
570
571 #[test]
572 fn test_string_similarity_near_match() {
573 let sim = EntityLinker::string_similarity("Einstein", "Einsten");
574 assert!(sim > 0.8, "sim = {}", sim);
575 }
576
577 #[test]
578 fn test_string_similarity_range() {
579 let sim = EntityLinker::string_similarity("kitten", "sitting");
580 assert!((0.0..=1.0).contains(&sim));
581 }
582
583 #[test]
586 fn test_edit_distance_identical() {
587 assert_eq!(EntityLinker::edit_distance("abc", "abc"), 0);
588 }
589
590 #[test]
591 fn test_edit_distance_one_empty() {
592 assert_eq!(EntityLinker::edit_distance("", "abc"), 3);
593 assert_eq!(EntityLinker::edit_distance("abc", ""), 3);
594 }
595
596 #[test]
597 fn test_edit_distance_kitten_sitting() {
598 assert_eq!(EntityLinker::edit_distance("kitten", "sitting"), 3);
599 }
600
601 #[test]
602 fn test_edit_distance_sunday_saturday() {
603 assert_eq!(EntityLinker::edit_distance("sunday", "saturday"), 3);
604 }
605
606 #[test]
607 fn test_edit_distance_single_char() {
608 assert_eq!(EntityLinker::edit_distance("a", "b"), 1);
609 assert_eq!(EntityLinker::edit_distance("a", "a"), 0);
610 }
611
612 #[test]
615 fn test_detect_mentions_exact_label() {
616 let linker = sample_linker();
617 let mentions = linker.detect_mentions("Albert Einstein was a physicist.");
618 assert!(!mentions.is_empty(), "expected at least one mention");
619 let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
620 assert!(
621 texts.iter().any(|t| t.to_lowercase().contains("einstein")),
622 "mentions = {:?}",
623 texts
624 );
625 }
626
627 #[test]
628 fn test_detect_mentions_alias() {
629 let linker = sample_linker();
630 let mentions = linker.detect_mentions("Curie discovered radium.");
631 let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
632 assert!(
633 texts.iter().any(|t| t.to_lowercase() == "curie"),
634 "mentions = {:?}",
635 texts
636 );
637 }
638
639 #[test]
640 fn test_detect_mentions_case_insensitive() {
641 let linker = sample_linker();
642 let mentions = linker.detect_mentions("berlin is a great city.");
643 assert!(!mentions.is_empty(), "expected mention of Berlin");
644 }
645
646 #[test]
647 fn test_detect_mentions_multiple() {
648 let linker = sample_linker();
649 let mentions = linker.detect_mentions("Einstein visited Berlin.");
650 assert!(mentions.len() >= 2, "mentions = {:?}", mentions.len());
652 }
653
654 #[test]
655 fn test_detect_mentions_no_match() {
656 let linker = sample_linker();
657 let mentions = linker.detect_mentions("The quick brown fox jumps.");
658 assert!(
659 mentions.is_empty(),
660 "expected no mentions, got {:?}",
661 mentions
662 );
663 }
664
665 #[test]
668 fn test_link_basic() {
669 let linker = sample_linker();
670 let linked = linker
671 .link("Albert Einstein won the Nobel Prize.")
672 .expect("ok");
673 assert!(!linked.is_empty());
674 }
675
676 #[test]
677 fn test_link_empty_text_error() {
678 let linker = sample_linker();
679 assert!(linker.link("").is_err());
680 assert!(linker.link(" ").is_err());
681 }
682
683 #[test]
684 fn test_link_empty_kb_error() {
685 let linker = EntityLinker::new(0.5);
686 assert!(matches!(
687 linker.link("some text"),
688 Err(LinkerError::EmptyKnowledgeBase)
689 ));
690 }
691
692 #[test]
693 fn test_link_confidence_populated() {
694 let linker = sample_linker();
695 let linked = linker.link("Curie was born in Poland.").expect("ok");
696 for le in &linked {
697 if le.best_candidate.is_some() {
698 assert!(le.confidence > 0.0);
699 }
700 }
701 }
702
703 #[test]
706 fn test_link_mention_exact() {
707 let linker = sample_linker();
708 let candidates = linker.link_mention("Einstein");
709 assert!(!candidates.is_empty());
710 assert_eq!(candidates[0].entity.iri, "http://example.org/einstein");
712 }
713
714 #[test]
715 fn test_link_mention_sorted_descending() {
716 let linker = sample_linker();
717 let candidates = linker.link_mention("Berlin");
718 for window in candidates.windows(2) {
719 assert!(
720 window[0].score >= window[1].score,
721 "candidates not sorted: {} < {}",
722 window[0].score,
723 window[1].score
724 );
725 }
726 }
727
728 #[test]
729 fn test_link_mention_returns_all_above_zero() {
730 let linker = sample_linker();
731 let candidates = linker.link_mention("Curie");
732 for c in &candidates {
734 assert!(c.score > 0.0);
735 }
736 }
737
738 #[test]
741 fn test_threshold_filters_low_confidence() {
742 let mut linker = EntityLinker::new(0.99); linker.add_entity(make_entity(
744 "http://x.org/z",
745 "Zephyr",
746 &[],
747 0.1, ));
749 let linked = linker.link("There is a zephyr wind.").expect("ok");
750 for le in &linked {
752 assert!(
753 le.best_candidate.is_none()
754 || le.best_candidate.as_ref().expect("some").score >= 0.99
755 );
756 }
757 }
758
759 #[test]
760 fn test_min_mention_len_filter() {
761 let mut linker = EntityLinker::new(0.0);
762 linker.add_entity(make_entity("http://x.org/a", "AI", &[], 0.5));
763 linker.min_mention_len = 5;
764 let mentions = linker.detect_mentions("AI is transforming the world.");
766 assert!(
767 mentions.is_empty() || mentions.iter().all(|m| m.text.len() >= 5),
768 "unexpected short mention found"
769 );
770 }
771
772 #[test]
775 fn test_linker_error_display() {
776 assert!(LinkerError::EmptyText.to_string().contains("empty"));
777 assert!(LinkerError::EmptyKnowledgeBase
778 .to_string()
779 .contains("no entities"));
780 assert!(LinkerError::InvalidEntity("bad".to_string())
781 .to_string()
782 .contains("bad"));
783 }
784
785 #[test]
788 fn test_combined_score_perfect() {
789 let linker = sample_linker();
791 let candidates = linker.link_mention("Albert Einstein");
792 if let Some(c) = candidates.first() {
794 assert!(c.score > 0.5, "score = {}", c.score);
795 }
796 }
797
798 #[test]
801 fn test_alias_detection_full_alias() {
802 let linker = sample_linker();
803 let mentions = linker.detect_mentions("A. Einstein changed physics.");
804 let texts: Vec<String> = mentions.iter().map(|m| m.text.to_lowercase()).collect();
805 assert!(
806 texts.iter().any(|t| t.contains("einstein")),
807 "texts = {:?}",
808 texts
809 );
810 }
811
812 #[test]
820 fn test_detect_mentions_indexed_lookup_with_many_entities() {
821 let mut linker = EntityLinker::new(0.3);
822 for i in 0..500 {
823 linker.add_entity(make_entity(
824 &format!("http://example.org/e{i}"),
825 &format!("Entity{i}"),
826 &[],
827 0.5,
828 ));
829 }
830 linker.add_entity(make_entity(
831 "http://example.org/einstein",
832 "Albert Einstein",
833 &["Einstein"],
834 0.95,
835 ));
836
837 let mentions = linker.detect_mentions("Albert Einstein was a physicist.");
838 let texts: Vec<&str> = mentions.iter().map(|m| m.text.as_str()).collect();
839 assert!(
840 texts.iter().any(|t| t.to_lowercase().contains("einstein")),
841 "mentions = {:?}",
842 texts
843 );
844 }
845
846 #[test]
850 fn test_link_mention_blocking_matches_correct_entity() {
851 let linker = sample_linker();
852 let candidates = linker.link_mention("Curie");
853 assert!(
854 candidates
855 .iter()
856 .any(|c| c.entity.iri == "http://example.org/curie"),
857 "candidates = {:?}",
858 candidates.iter().map(|c| &c.entity.iri).collect::<Vec<_>>()
859 );
860 }
861
862 #[test]
865 fn test_link_mention_blocking_no_match_for_unindexed_first_char() {
866 let linker = sample_linker();
867 let candidates = linker.link_mention("Zzzzz nonexistent name");
868 assert!(candidates.is_empty(), "candidates = {:?}", candidates);
869 }
870}