use crate::knowledge::{EntityId, EntityKey, EntityTypeRef, ResolutionMethod};
use oxibrain_index::ngram;
use std::collections::BTreeMap;
#[derive(Debug, Clone)]
pub struct PerType<T: Clone> {
default: T,
overrides: BTreeMap<String, T>,
}
impl<T: Clone> PerType<T> {
pub fn new(default: T) -> Self {
Self {
default,
overrides: BTreeMap::new(),
}
}
pub fn set(&mut self, ty: &str, value: T) {
self.overrides.insert(ty.to_string(), value);
}
pub fn get(&self, ty: &str) -> &T {
self.overrides.get(ty).unwrap_or(&self.default)
}
}
impl PerType<f64> {
pub fn weight(&self, ty: &str) -> f64 {
*self.get(ty)
}
}
#[derive(Debug, Clone)]
pub struct ResolutionConfig {
pub tau_high: f64,
pub tau_low: f64,
pub w_exact: f64,
pub w_ngram: f64,
pub w_graph: f64,
pub w_embedding: PerType<f64>,
}
impl Default for ResolutionConfig {
fn default() -> Self {
Self {
tau_high: 0.75,
tau_low: 0.25,
w_exact: 1.0,
w_ngram: 1.0,
w_graph: 0.4,
w_embedding: {
let mut w = PerType::new(0.3);
w.set("Person", 0.1);
w.set("Organization", 0.1);
w.set("Concept", 0.6);
w
},
}
}
}
#[derive(Debug, Clone)]
pub enum Decision {
Link {
entity: EntityId,
method: ResolutionMethod,
score: f64,
},
New {
method: ResolutionMethod,
score: f64,
},
Candidate {
new_entity: EntityId,
existing: EntityId,
score: f64,
},
}
pub fn normalize(surface: &str, _ty: &EntityTypeRef) -> String {
use unicode_normalization::UnicodeNormalization;
surface
.nfkc()
.collect::<String>()
.to_lowercase()
.split_whitespace()
.collect::<Vec<_>>()
.join(" ")
}
pub fn score(
candidate: &EntityKey,
mention_normalized: &str,
mention_type: &EntityTypeRef,
graph_context: f64,
embedding_sim: f64,
config: &ResolutionConfig,
) -> f64 {
if candidate.ty != *mention_type {
return 0.0;
}
let exact = if candidate.normalized == mention_normalized {
1.0
} else {
0.0
};
let cand_shingles = ngram::shingles(&candidate.normalized, 3);
let ment_shingles = ngram::shingles(mention_normalized, 3);
let j = ngram::jaccard(&cand_shingles, &ment_shingles);
let emb_weight = config.w_embedding.weight(mention_type);
let raw = config.w_exact * exact
+ config.w_ngram * j
+ config.w_graph * graph_context
+ emb_weight * embedding_sim;
raw.clamp(0.0, 1.0)
}
pub fn resolve(
mention_normalized: &str,
mention_type: &EntityTypeRef,
candidates: &[EntityKey],
graph_context: impl Fn(&EntityId) -> f64,
embedding_sim: impl Fn(&EntityId) -> f64,
config: &ResolutionConfig,
) -> Decision {
let mut scored: Vec<(f64, &EntityKey)> = Vec::new();
for c in candidates {
let ctx = graph_context(&c.entity);
let emb = embedding_sim(&c.entity);
let s = score(c, mention_normalized, mention_type, ctx, emb, config);
if s > 0.0 {
scored.push((s, c));
}
}
scored.sort_by(|a, b| {
b.0.partial_cmp(&a.0)
.unwrap_or(std::cmp::Ordering::Equal)
.then(a.1.entity.cmp(&b.1.entity))
});
match scored.first() {
None => Decision::New {
method: ResolutionMethod::New,
score: 0.0,
},
Some(&(best, c)) if best >= config.tau_high => {
let method = if c.normalized == mention_normalized {
ResolutionMethod::ExactKey
} else {
ResolutionMethod::Lexical { score: best }
};
Decision::Link {
entity: c.entity.clone(),
method,
score: best,
}
}
Some(&(best, _c)) if best <= config.tau_low => Decision::New {
method: ResolutionMethod::New,
score: best,
},
Some(&(best, c)) => {
Decision::Candidate {
new_entity: String::new(), existing: c.entity.clone(),
score: best,
}
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::knowledge::KeyOrigin;
fn make_key(entity: &str, normalized: &str, ty: &str) -> EntityKey {
EntityKey {
id: format!("k_{entity}_{normalized}"),
space: "s1".into(),
entity: entity.into(),
ty: ty.into(),
normalized: normalized.into(),
surface: normalized.into(),
origin: KeyOrigin::UserDeclared,
}
}
#[test]
fn exact_match_links() {
let cands = vec![make_key("e1", "alice", "Person")];
let dec = resolve(
"alice",
&"Person".to_string(),
&cands,
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
match dec {
Decision::Link {
entity,
method,
score,
} => {
assert_eq!(entity, "e1");
assert!(score >= 0.75);
assert!(matches!(method, ResolutionMethod::ExactKey));
}
_ => panic!("expected Link"),
}
}
#[test]
fn type_mismatch_rejected() {
let cands = vec![make_key("e1", "alice", "Organization")];
let dec = resolve(
"alice",
&"Person".to_string(),
&cands,
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
assert!(matches!(dec, Decision::New { .. }));
}
#[test]
fn no_candidates_is_new() {
let dec = resolve(
"alice",
&"Person".to_string(),
&[],
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
assert!(matches!(dec, Decision::New { .. }));
}
#[test]
fn normalize_basic() {
assert_eq!(normalize("Alice", &"Person".to_string()), "alice");
assert_eq!(
normalize(" Alice Smith ", &"Person".to_string()),
"alice smith"
);
}
#[test]
fn low_similarity_is_new() {
let cands = vec![make_key("e1", "zzzzzzzzz", "Person")];
let dec = resolve(
"alice",
&"Person".to_string(),
&cands,
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
assert!(matches!(dec, Decision::New { .. }));
}
#[test]
fn near_match_without_context_is_candidate() {
let cands = vec![make_key("e1", "alicia", "Person")];
let dec = resolve(
"alice",
&"Person".to_string(),
&cands,
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
assert!(
matches!(dec, Decision::Candidate { .. }),
"near match without context should be Candidate, got {dec:?}"
);
}
#[test]
fn near_match_with_context_links() {
let cands = vec![make_key("e1", "alicia", "Person")];
let dec = resolve(
"alice",
&"Person".to_string(),
&cands,
|_| 1.0, |_| 0.0,
&ResolutionConfig::default(),
);
assert!(
matches!(dec, Decision::Link { .. }),
"near match with context should Link, got {dec:?}"
);
}
#[test]
fn prefix_sharing_does_not_inflate_score() {
let cands = vec![make_key("e1", "김서연", "Person")];
let dec = resolve(
"김민수",
&"Person".to_string(),
&cands,
|_| 0.0,
|_| 0.0,
&ResolutionConfig::default(),
);
assert!(
!matches!(dec, Decision::Link { .. }),
"shared surname should not Link without context, got {dec:?}"
);
}
#[test]
fn pertype_default_and_override() {
let mut pt = PerType::new(0.0);
assert_eq!(pt.weight("Person"), 0.0);
pt.set("Concept", 0.3);
assert_eq!(pt.weight("Concept"), 0.3);
assert_eq!(pt.weight("Person"), 0.0); }
}