Spaces:
Sleeping
Sleeping
File size: 6,899 Bytes
db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f db4d559 3786a3f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | import logging
from typing import List, Dict, Optional
from rapidfuzz import fuzz
logger = logging.getLogger(__name__)
class EntityResolver:
def __init__(self, similarity_threshold: float = 85.0):
self.similarity_threshold = similarity_threshold
self.llm = None
self._init_llm()
logger.info("Initializing EntityResolver with similarity threshold: %.1f", self.similarity_threshold)
def _init_llm(self):
"""Initialize LLM for borderline disambiguation. Caches the instance."""
try:
from .llm_client import get_llm
self.llm = get_llm(temperature=0.0)
except Exception as e:
logger.warning("Could not initialize LLM for entity disambiguation: %s", str(e))
self.llm = None
def resolve_entities(self, entities: List[dict], relationships: List[dict]) -> tuple[List[dict], List[dict]]:
"""
Deduplicates a list of entities and rewrites relationship references accordingly.
Returns a tuple of (resolved_entities, rewritten_relationships).
"""
if not entities:
return [], relationships
logger.info("Resolving duplicates for %d entities and %d relationships...", len(entities), len(relationships))
# 1. Group entities by their category/type
grouped_by_type: Dict[str, List[dict]] = {}
for entity in entities:
etype = entity['type'].upper()
grouped_by_type.setdefault(etype, []).append(entity)
resolved_entities = []
name_mappings = {} # Maps original_name -> canonical_name
for etype, group in grouped_by_type.items():
resolved_group = []
for current in group:
matched_canonical = None
# Check current entity against already resolved entities in the same type group
for existing in resolved_group:
# Run fuzzy comparison on names
ratio = fuzz.token_set_ratio(current['name'].lower(), existing['name'].lower())
if ratio >= self.similarity_threshold:
matched_canonical = existing
break
# Borderline cases: ask LLM to disambiguate
elif self.similarity_threshold - 15 <= ratio < self.similarity_threshold:
if self._llm_confirm_duplicate(
current['name'], existing['name'],
current.get('type', ''), existing.get('type', '')
):
matched_canonical = existing
break
if matched_canonical:
# Duplicate found! Merge current into matched_canonical
old_name = current['name']
new_name = matched_canonical['name']
# Keep the longer name as the canonical one
if len(old_name) > len(new_name):
matched_canonical['name'] = old_name
name_mappings[new_name] = old_name
name_mappings[old_name] = old_name
else:
name_mappings[old_name] = new_name
# Combine descriptions, avoiding exact duplicates
if current['description'] and current['description'] not in matched_canonical['description']:
matched_canonical['description'] = f"{matched_canonical['description']} {current['description']}".strip()
logger.info("Resolved & Merged entity: '%s' ➔ '%s'", old_name, matched_canonical['name'])
else:
# Unique entity in this pass, add it to resolved list
resolved_group.append(current)
name_mappings[current['name']] = current['name']
resolved_entities.extend(resolved_group)
# 2. Rewrite relationships using the canonical name mappings
rewritten_relationships = []
for rel in relationships:
# Look up source and target names in mappings
src = rel['source_entity']
tgt = rel['target_entity']
canonical_src = name_mappings.get(src, src)
canonical_tgt = name_mappings.get(tgt, tgt)
# Prevent self-referencing relationships created by merges
if canonical_src == canonical_tgt:
logger.warning("Discarded self-referencing relationship: [%s] --[%s]--> [%s] after resolution merge.",
src, rel['relationship_type'], tgt)
continue
rel['source_entity'] = canonical_src
rel['target_entity'] = canonical_tgt
rewritten_relationships.append(rel)
logger.info("Deduplication complete. Resolved entities count: %d (from %d) | Relationships count: %d",
len(resolved_entities), len(entities), len(rewritten_relationships))
return resolved_entities, rewritten_relationships
def _llm_confirm_duplicate(self, name_a: str, name_b: str, type_a: str, type_b: str) -> bool:
"""
Uses LLM to determine if two borderline entities refer to the same real-world thing.
Returns True if they are the same entity.
"""
if not self.llm:
return False
try:
from langchain_core.prompts import ChatPromptTemplate
prompt = ChatPromptTemplate.from_messages([
("system", (
"You are an expert at entity resolution. Given two entity names and their types, "
"determine if they refer to the same real-world entity.\n"
"Answer ONLY 'yes' or 'no'.\n"
"Consider: abbreviations, nicknames, partial names, and variations of the same entity."
)),
("human", (
"Entity A: '{name_a}' (Type: {type_a})\n"
"Entity B: '{name_b}' (Type: {type_b})\n\n"
"Do these refer to the same entity?"
))
])
chain = prompt | self.llm
response = chain.invoke({
"name_a": name_a,
"name_b": name_b,
"type_a": type_a,
"type_b": type_b
})
answer = response.content.strip().lower()
is_duplicate = answer.startswith("yes")
logger.info("LLM disambiguation: '%s' vs '%s' -> %s", name_a, name_b, is_duplicate)
return is_duplicate
except Exception as e:
logger.error("LLM disambiguation failed: %s", str(e))
return False
|