Upload notebooks/02_entity_network/21_entity_resolution.ipynb with huggingface_hub

notebooks/02_entity_network/21_entity_resolution.ipynb

@@ -0,0 +1,302 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# 21 - Entity Resolution\n",
+    "\n",
+    "Pipeline notebook for fuzzy entity matching and alias resolution.\n",
+    "\n",
+    "Groups similar entity names into clusters using normalized forms and fuzzy string matching.\n",
+    "For each cluster, picks the most frequent form as the canonical name.\n",
+    "Results are stored in the `entity_aliases` table."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "tags": [
+     "parameters"
+    ]
+   },
+   "outputs": [],
+   "source": [
+    "# Parameters\n",
+    "source_section = None\n",
+    "similarity_threshold = 0.85"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "sys.path.insert(0, '/opt/epstein_env/research')\n",
+    "\n",
+    "import difflib\n",
+    "from collections import defaultdict, Counter\n",
+    "from tqdm.auto import tqdm\n",
+    "\n",
+    "from research_lib.db import fetch_df, fetch_all, bulk_insert\n",
+    "from research_lib.nlp import normalize_entity\n",
+    "from research_lib.incremental import start_run, finish_run"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Start run\n",
+    "run_id = start_run(\n",
+    "    'entity_resolution',\n",
+    "    source_section=source_section,\n",
+    "    parameters={'similarity_threshold': similarity_threshold},\n",
+    ")\n",
+    "print(f'Started run {run_id}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Load all unique entity texts with their frequencies, grouped by type\n",
+    "where_clause = ''\n",
+    "params = []\n",
+    "if source_section:\n",
+    "    where_clause = 'WHERE d.source_section = %s'\n",
+    "    params = [source_section]\n",
+    "\n",
+    "sql = f\"\"\"\n",
+    "    SELECT e.entity_text, e.entity_type, COUNT(*) as freq\n",
+    "    FROM entities e\n",
+    "    JOIN documents d ON d.id = e.document_id\n",
+    "    {where_clause}\n",
+    "    GROUP BY e.entity_text, e.entity_type\n",
+    "    ORDER BY freq DESC\n",
+    "\"\"\"\n",
+    "entity_df = fetch_df(sql, params or None)\n",
+    "print(f'Total unique entity-type combinations: {len(entity_df)}')\n",
+    "print(f'Entity types: {entity_df[\"entity_type\"].value_counts().to_dict()}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Normalize entity texts\n",
+    "entity_df['normalized'] = entity_df.apply(\n",
+    "    lambda row: normalize_entity(row['entity_text'], row['entity_type']),\n",
+    "    axis=1,\n",
+    ")\n",
+    "\n",
+    "# Group entities by type for within-type matching\n",
+    "entities_by_type = {}\n",
+    "for etype, group in entity_df.groupby('entity_type'):\n",
+    "    entities_by_type[etype] = group.reset_index(drop=True)\n",
+    "    print(f'  {etype}: {len(group)} unique entities')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def find_clusters(entities_group, threshold):\n",
+    "    \"\"\"Cluster entities using fuzzy matching with Union-Find.\"\"\"\n",
+    "    texts = entities_group['normalized'].tolist()\n",
+    "    freqs = entities_group['freq'].tolist()\n",
+    "    originals = entities_group['entity_text'].tolist()\n",
+    "    n = len(texts)\n",
+    "\n",
+    "    # Union-Find\n",
+    "    parent = list(range(n))\n",
+    "\n",
+    "    def find(x):\n",
+    "        while parent[x] != x:\n",
+    "            parent[x] = parent[parent[x]]\n",
+    "            x = parent[x]\n",
+    "        return x\n",
+    "\n",
+    "    def union(a, b):\n",
+    "        ra, rb = find(a), find(b)\n",
+    "        if ra != rb:\n",
+    "            # Attach less frequent to more frequent\n",
+    "            if freqs[ra] >= freqs[rb]:\n",
+    "                parent[rb] = ra\n",
+    "            else:\n",
+    "                parent[ra] = rb\n",
+    "\n",
+    "    # Compare all pairs using SequenceMatcher\n",
+    "    # For efficiency, first group by exact normalized form\n",
+    "    norm_groups = defaultdict(list)\n",
+    "    for i, norm in enumerate(texts):\n",
+    "        norm_groups[norm.lower()].append(i)\n",
+    "\n",
+    "    # Union exact normalized matches\n",
+    "    for indices in norm_groups.values():\n",
+    "        for j in range(1, len(indices)):\n",
+    "            union(indices[0], indices[j])\n",
+    "\n",
+    "    # Fuzzy match across distinct normalized forms\n",
+    "    unique_norms = list(norm_groups.keys())\n",
+    "    for i in tqdm(range(len(unique_norms)), desc='Fuzzy matching', leave=False):\n",
+    "        for j in range(i + 1, len(unique_norms)):\n",
+    "            ratio = difflib.SequenceMatcher(\n",
+    "                None, unique_norms[i], unique_norms[j]\n",
+    "            ).ratio()\n",
+    "            if ratio >= threshold:\n",
+    "                # Union representatives from each group\n",
+    "                union(norm_groups[unique_norms[i]][0], norm_groups[unique_norms[j]][0])\n",
+    "\n",
+    "    # Build clusters\n",
+    "    clusters = defaultdict(list)\n",
+    "    for i in range(n):\n",
+    "        root = find(i)\n",
+    "        clusters[root].append(i)\n",
+    "\n",
+    "    # Pick canonical name (most frequent original form)\n",
+    "    result = []\n",
+    "    for root, members in clusters.items():\n",
+    "        if len(members) <= 1:\n",
+    "            continue  # Skip singletons\n",
+    "        best_idx = max(members, key=lambda i: freqs[i])\n",
+    "        canonical = originals[best_idx]\n",
+    "        for idx in members:\n",
+    "            if idx != best_idx:\n",
+    "                result.append({\n",
+    "                    'alias_text': originals[idx],\n",
+    "                    'canonical_text': canonical,\n",
+    "                    'similarity': difflib.SequenceMatcher(\n",
+    "                        None,\n",
+    "                        texts[idx].lower(),\n",
+    "                        texts[best_idx].lower(),\n",
+    "                    ).ratio(),\n",
+    "                })\n",
+    "\n",
+    "    return result, len(clusters)\n",
+    "\n",
+    "print('Cluster function defined.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Run clustering for each entity type\n",
+    "all_aliases = []\n",
+    "total_clusters = 0\n",
+    "total_entities = 0\n",
+    "\n",
+    "for etype, group in entities_by_type.items():\n",
+    "    print(f'\\nProcessing {etype} ({len(group)} entities)...')\n",
+    "    total_entities += len(group)\n",
+    "\n",
+    "    aliases, n_clusters = find_clusters(group, similarity_threshold)\n",
+    "    total_clusters += n_clusters\n",
+    "\n",
+    "    for alias in aliases:\n",
+    "        alias['entity_type'] = etype\n",
+    "\n",
+    "    all_aliases.extend(aliases)\n",
+    "    print(f'  Found {n_clusters} clusters, {len(aliases)} alias mappings')\n",
+    "\n",
+    "print(f'\\nTotal aliases found: {len(all_aliases)}')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Insert into entity_aliases table\n",
+    "rows = [\n",
+    "    (\n",
+    "        a['alias_text'],\n",
+    "        a['canonical_text'],\n",
+    "        a['entity_type'],\n",
+    "        a['similarity'],\n",
+    "        source_section,\n",
+    "    )\n",
+    "    for a in all_aliases\n",
+    "]\n",
+    "\n",
+    "if rows:\n",
+    "    inserted = bulk_insert(\n",
+    "        'entity_aliases',\n",
+    "        ['alias_text', 'canonical_text', 'entity_type', 'similarity_score', 'source_section'],\n",
+    "        rows,\n",
+    "        on_conflict='DO NOTHING',\n",
+    "    )\n",
+    "    print(f'Inserted {inserted} alias rows')\n",
+    "else:\n",
+    "    print('No aliases to insert.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Finish run\n",
+    "finish_run(run_id, documents_processed=total_entities)\n",
+    "print(f'Run {run_id} completed.')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Summary stats\n",
+    "print('=== Entity Resolution Summary ===')\n",
+    "print(f'Total unique entities: {total_entities}')\n",
+    "print(f'Total clusters (multi-member): {total_clusters}')\n",
+    "print(f'Total alias mappings: {len(all_aliases)}')\n",
+    "if total_entities > 0:\n",
+    "    reduction = len(all_aliases) / total_entities * 100\n",
+    "    print(f'Reduction ratio: {reduction:.1f}% of entities are aliases')\n",
+    "\n",
+    "# Show some example clusters\n",
+    "if all_aliases:\n",
+    "    from collections import defaultdict\n",
+    "    clusters_display = defaultdict(list)\n",
+    "    for a in all_aliases[:100]:\n",
+    "        clusters_display[a['canonical_text']].append(\n",
+    "            f\"{a['alias_text']} ({a['similarity']:.2f})\"\n",
+    "        )\n",
+    "    print('\\nExample clusters (first 10):')\n",
+    "    for i, (canonical, aliases) in enumerate(list(clusters_display.items())[:10]):\n",
+    "        print(f'  {canonical}: {aliases}')"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}