CoolFace
Apppublic

GodsDevProject/FOIA_Doc_Search

sourceHugging Facemitupdated 8mo agoView on Hugging Face
1likes
entity_graph.py19 linesDownload Raw Back to root
1import networkx as nx2from typing import List, Dict3 4def build_entity_graph(docs: List[Dict]) -> Dict:5    G = nx.Graph()6 7    for d in docs:8        agency = d.get("agency", "Unknown")9        G.add_node(agency, group="agency")10 11        for token in d.get("content", "").split():12            if token.isupper() and len(token) > 2:13                G.add_node(token, group="entity")14                G.add_edge(agency, token)15 16    return {17        "nodes": [{"id": n, "group": G.nodes[n]["group"]} for n in G.nodes],18        "links": [{"source": u, "target": v} for u, v in G.edges]19    }