SuperScraper/entity_resolver.py at main · amish-gaur/SuperScraper · GitHub

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
"""Entity normalization and fuzzy-resolution helpers for dataset assembly."""

from __future__ import annotations

from dataclasses import dataclass
from difflib import SequenceMatcher
import re

import pandas as pd


COMMON_ENTITY_SUFFIXES = {
    "inc",
    "incorporated",
    "corp",
    "corporation",
    "company",
    "co",
    "ltd",
    "llc",
    "plc",
    "group",
    "holdings",
}


@dataclass(slots=True)
class EntityResolver:
    """Normalize entity labels and collapse near-duplicate keys."""

    similarity_threshold: float = 0.94

    def canonical_key(self, value: str | object) -> str:
        if value is None or pd.isna(value):
            return ""
        value = str(value).strip()
        if not value or value.lower() in {"nan", "none"}:
            return ""
        compact_ticker_pattern = bool(
            re.search(r"[a-z][A-Z]{2,}$", value)
            or re.search(r"\d{3,5}\.[A-Z]{1,3}$", value)
            or re.search(r"[A-Za-z]+\d{2,5}$", value)
        )
        value = re.sub(r"([a-z])([A-Z][a-z])", r"\1 \2", value)
        value = re.sub(r"([A-Za-z])(\d)", r"\1 \2", value)
        value = re.sub(r"(\d)([A-Za-z])", r"\1 \2", value)
        lowered = value.lower()
        lowered = re.sub(r"\[[^\]]+\]", "", lowered)
        lowered = lowered.replace("&", " and ")
        lowered = re.sub(r"[^a-z0-9\s]", " ", lowered)
        tokens = [token for token in lowered.split() if token and token not in COMMON_ENTITY_SUFFIXES]
        if len(tokens) >= 2 and tokens[-1].isalpha() and len(tokens[-1]) <= 2:
            tokens = tokens[:-1]
        if compact_ticker_pattern:
            if len(tokens) >= 2 and tokens[-1].isalpha() and len(tokens[-1]) <= 5:
                tokens = tokens[:-1]
            if len(tokens) >= 2 and tokens[-1].isdigit() and len(tokens[-1]) <= 5:
                tokens = tokens[:-1]
            if len(tokens) >= 2 and tokens[-1].isalpha() and len(tokens[-1]) <= 3:
                tokens = tokens[:-1]
        return " ".join(tokens)

    def resolve_frame(self, frame: pd.DataFrame, *, entity_column: str) -> pd.DataFrame:
        """Rewrite an entity column so near duplicates collapse to one canonical key."""
        resolved = frame.copy()
        canonical_map: dict[str, str] = {}
        seen_keys: list[str] = []

        for raw_value in resolved[entity_column].astype(str):
            key = self.canonical_key(raw_value)
            matched = self._best_match(key, seen_keys)
            canonical = matched or key
            canonical_map[raw_value] = canonical
            if not matched:
                seen_keys.append(canonical)

        resolved[entity_column] = resolved[entity_column].astype(str).map(canonical_map)
        return resolved

    def _best_match(self, key: str, candidates: list[str]) -> str | None:
        for candidate in candidates:
            if key == candidate:
                return candidate
            if SequenceMatcher(a=key, b=candidate).ratio() >= self.similarity_threshold:
                return candidate
        return None