Spaces:
Running on Zero
Running on Zero
| """Entity linking from question text (or answer strings) to graph nodes. | |
| Order of evidence: (1) exact match of a node name or alias as a whole word/phrase, | |
| longest first, case-insensitive; (2) paper titles by fuzzy match (questions quote | |
| titles); (3) rapidfuzz ratio of normalised strings for single answer strings. | |
| """ | |
| from __future__ import annotations | |
| import re | |
| from dataclasses import dataclass | |
| from rapidfuzz import fuzz, process | |
| from .canonical import normalise | |
| from .graph import KG | |
| class Link: | |
| node_id: str | |
| surface: str | |
| score: float | |
| method: str | |
| class EntityLinker: | |
| def __init__(self, kg: KG, fuzzy_threshold: float = 90.0, min_alias_len: int = 2): | |
| self.kg = kg | |
| self.fuzzy_threshold = fuzzy_threshold | |
| self.alias_index: dict[str, set[str]] = {} | |
| for n in kg.nodes.values(): | |
| if n.type == "Paper": | |
| continue | |
| # The canonical key itself (e.g. an abbreviation) may never occur as a surface form. | |
| for form in {n.name, *n.aliases, n.node_id.split(":", 1)[1]}: | |
| key = normalise(form, n.type) | |
| if len(key) >= min_alias_len: | |
| self.alias_index.setdefault(key, set()).add(n.node_id) | |
| raw = normalise(form) | |
| if len(raw) >= min_alias_len: | |
| self.alias_index.setdefault(raw, set()).add(n.node_id) | |
| self.titles = {normalise(n.name): n.node_id for n in kg.nodes.values() if n.type == "Paper"} | |
| alternation = "|".join(re.escape(k) for k in sorted(self.alias_index, key=len, reverse=True)) | |
| self._pattern = re.compile(rf"(?<![a-z0-9])(?:{alternation})(?![a-z0-9])") if alternation else None | |
| def link_question(self, question: str) -> list[Link]: | |
| q = normalise(question) | |
| links: list[Link] = [] | |
| taken: list[tuple[int, int]] = [] | |
| # Paper titles first: a quoted title often contains method/dataset names. | |
| for m in re.finditer(r'"([^"]{8,})"', question): | |
| hit = self.match_title(m.group(1)) | |
| if hit: | |
| links.append(hit) | |
| tq = normalise(m.group(1)) | |
| p = q.find(tq) | |
| if p >= 0: | |
| taken.append((p, p + len(tq))) | |
| if self._pattern is not None: | |
| for m in self._pattern.finditer(q): | |
| if any(a <= m.start() < b for a, b in taken): | |
| continue | |
| for nid in sorted(self.alias_index[m.group(0)]): | |
| links.append(Link(nid, m.group(0), 100.0, "exact")) | |
| seen: set[str] = set() | |
| out = [] | |
| for lk in links: | |
| if lk.node_id not in seen: | |
| seen.add(lk.node_id) | |
| out.append(lk) | |
| return out | |
| def match_title(self, text: str) -> Link | None: | |
| if not self.titles: | |
| return None | |
| best = process.extractOne(normalise(text), list(self.titles), scorer=fuzz.ratio) | |
| if best and best[1] >= self.fuzzy_threshold: | |
| return Link(self.titles[best[0]], text, float(best[1]), "title") | |
| return None | |
| def link_name(self, name: str, etype: str | None = None) -> Link | None: | |
| """Link a single entity string (an answer item or a path element).""" | |
| key_typed = normalise(name, etype) if etype else None | |
| for key in filter(None, (key_typed, normalise(name), normalise(name, "Dataset"))): | |
| ids = sorted(self.alias_index.get(key, ())) | |
| if etype: | |
| ids = [i for i in ids if self.kg.nodes[i].type == etype] or ids | |
| if ids: | |
| return Link(ids[0], name, 100.0, "exact") | |
| title = self.match_title(name) | |
| if title: | |
| return title | |
| if not self.alias_index: | |
| return None | |
| best = process.extractOne(normalise(name), list(self.alias_index), scorer=fuzz.ratio) | |
| if best and best[1] >= self.fuzzy_threshold: | |
| return Link(min(self.alias_index[best[0]]), name, float(best[1]), "fuzzy") | |
| return None | |