Download src/alignment.py from eduzrh/demo: direct link, hf CLI and curl.
- Browser
- Download file 11.6 kB
-
https://huggingface.co/spaces/eduzrh/demo/resolve/main/src/alignment.py
- Command line
-
hf download hf://spaces/eduzrh/demo/src/alignment.py
-
curl -L -o alignment.py https://huggingface.co/spaces/eduzrh/demo/resolve/main/src/alignment.py
11.6 kB
| """ | |
| AgentEA-Studio: Entity Alignment Pipeline (lightweight). | |
| Implements the AdaCoAgentEA multi-agent alignment pipeline. | |
| Uses TF-IDF + NumPy for embeddings (no heavy ML deps). | |
| """ | |
| import time | |
| import numpy as np | |
| from collections import defaultdict | |
| from sklearn.feature_extraction.text import TfidfVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| class AlignmentPipeline: | |
| """Multi-agent entity alignment pipeline (lightweight).""" | |
| def __init__(self, kg1, kg2, verbose=True): | |
| self.kg1 = kg1 | |
| self.kg2 = kg2 | |
| self.verbose = verbose | |
| self.logs = [] | |
| self.results = {} | |
| # Initialize TF-IDF vectorizer for S1 | |
| self._log("Initializing TF-IDF vectorizer...") | |
| all_names = [e["name"] for e in kg1.entities.values()] + \ | |
| [e["name"] for e in kg2.entities.values()] | |
| self.vectorizer = TfidfVectorizer(analyzer="char_wb", ngram_range=(2, 4), max_features=256) | |
| self.vectorizer.fit(all_names) | |
| self._log("TF-IDF vectorizer ready.") | |
| def _log(self, msg): | |
| if self.verbose: | |
| print(msg) | |
| self.logs.append({"time": time.time(), "message": msg}) | |
| def run_full_pipeline(self): | |
| """Run the complete multi-agent alignment pipeline.""" | |
| self._log("=" * 60) | |
| self._log("Starting AdaCoAgentEA Multi-Agent Alignment Pipeline") | |
| self._log("=" * 60) | |
| # === Area 1: Structural-Semantic Decoupling === | |
| self._log("\n[Area 1] Structural-Semantic Decoupling") | |
| # S1: Candidate Entity Retrieval | |
| candidates = self._run_s1_candidate_retrieval() | |
| self.results["s1_candidates"] = candidates | |
| # S2: Key Node Analysis | |
| key_nodes = self._run_s2_key_node_analysis() | |
| self.results["s2_key_nodes"] = key_nodes | |
| # S3: Extreme Information Density Parsing | |
| extreme_entities = self._run_s3_extreme_density_parsing() | |
| self.results["s3_extreme_entities"] = extreme_entities | |
| # === Gating G1 === | |
| g1 = self._evaluate_gating_g1() | |
| self._log(f"Gating G1: {'Activated' if g1 else 'Suppressed'}") | |
| # === Area 2: Deep Localization === | |
| self._log("\n[Area 2] Deep Localization") | |
| # L1: Heterogeneous Entity Selection | |
| high_conf_pairs = self._run_l1_high_conf_selection(candidates) | |
| self.results["l1_high_conf_pairs"] = high_conf_pairs | |
| # S4: Large-Scale Accurate Alignment | |
| similarity_matrix, uncertain_pairs = self._run_s4_large_scale_alignment(high_conf_pairs) | |
| self.results["s4_similarity_matrix"] = similarity_matrix | |
| self.results["s4_uncertain_pairs"] = uncertain_pairs | |
| # L2: Multi-Granularity Summarization | |
| meta_rules = self._run_l2_meta_rules(high_conf_pairs, uncertain_pairs) | |
| self.results["l2_meta_rules"] = meta_rules | |
| # === Gating === | |
| g3 = len(uncertain_pairs) > 0 | |
| self._log(f"Gating G3: {'Activated' if g3 else 'Suppressed'}") | |
| # === Area 3: Knowledge Aggregation === | |
| if g3: | |
| self._log("\n[Area 3] Knowledge Aggregation") | |
| self._run_l3_l4_knowledge_aggregation(uncertain_pairs, meta_rules) | |
| # === Final Results === | |
| self._log("\n" + "=" * 60) | |
| self._log("Pipeline Complete!") | |
| self._log(f" High-confidence pairs: {len(high_conf_pairs)}") | |
| self._log(f" Uncertain pairs: {len(uncertain_pairs)}") | |
| self._log(f" Meta-logic rules: {len(meta_rules)}") | |
| self._log("=" * 60) | |
| return { | |
| "high_conf_pairs": high_conf_pairs, | |
| "uncertain_pairs": uncertain_pairs, | |
| "similarity_matrix": similarity_matrix, | |
| "meta_rules": meta_rules, | |
| "key_nodes": key_nodes, | |
| "extreme_entities": extreme_entities, | |
| "logs": self.logs, | |
| } | |
| def _run_s1_candidate_retrieval(self): | |
| """S1: Retrieve top-k candidate entities using TF-IDF similarity.""" | |
| self._log("[S1] Candidate Entity Retrieval — building TF-IDF index...") | |
| eids_kg2 = list(self.kg2.entities.keys()) | |
| names_kg2 = [self.kg2.entities[eid]["name"] for eid in eids_kg2] | |
| emb_kg2 = self.vectorizer.transform(names_kg2).toarray() | |
| candidate_map = {} | |
| for eid1, e1 in self.kg1.entities.items(): | |
| emb1 = self.vectorizer.transform([e1["name"]]).toarray() | |
| sims = cosine_similarity(emb1, emb_kg2)[0] | |
| top_k = min(5, len(eids_kg2)) | |
| sorted_indices = np.argsort(sims)[::-1] | |
| seen = set() | |
| top_indices = [] | |
| for idx in sorted_indices: | |
| eid2 = eids_kg2[idx] | |
| if eid2 not in seen: | |
| seen.add(eid2) | |
| top_indices.append(idx) | |
| if len(top_indices) >= top_k: | |
| break | |
| candidates = [] | |
| for idx in top_indices: | |
| eid2 = eids_kg2[idx] | |
| candidates.append({ | |
| "eid": eid2, | |
| "name": self.kg2.entities[eid2]["name"], | |
| "score": float(sims[idx]), | |
| }) | |
| candidate_map[eid1] = candidates | |
| self._log(f" Retrieved candidates for {len(candidate_map)} entities") | |
| return candidate_map | |
| def _run_s2_key_node_analysis(self): | |
| """S2: Identify key nodes using PageRank.""" | |
| self._log("[S2] Key Node Analysis — running PageRank...") | |
| key_nodes_kg1 = self.kg1.get_key_nodes(method="pagerank", top_k=15) | |
| key_nodes_kg2 = self.kg2.get_key_nodes(method="pagerank", top_k=15) | |
| self._log(f" KG1: {len(key_nodes_kg1)} key nodes, KG2: {len(key_nodes_kg2)} key nodes") | |
| return {"kg1": key_nodes_kg1, "kg2": key_nodes_kg2} | |
| def _run_s3_extreme_density_parsing(self): | |
| """S3: Identify entities with extreme information density.""" | |
| self._log("[S3] Extreme Information Density Parsing...") | |
| high_deg_kg1 = self.kg1.get_high_degree_entities(threshold_percentile=90) | |
| high_deg_kg2 = self.kg2.get_high_degree_entities(threshold_percentile=90) | |
| low_deg_kg1 = self.kg1.get_low_degree_entities(max_degree=2) | |
| low_deg_kg2 = self.kg2.get_low_degree_entities(max_degree=2) | |
| self._log(f" KG1: {len(high_deg_kg1)} high-degree, {len(low_deg_kg1)} low-degree") | |
| self._log(f" KG2: {len(high_deg_kg2)} high-degree, {len(low_deg_kg2)} low-degree") | |
| return { | |
| "kg1": {"high_degree": high_deg_kg1, "low_degree": low_deg_kg1}, | |
| "kg2": {"high_degree": high_deg_kg2, "low_degree": low_deg_kg2}, | |
| } | |
| def _evaluate_gating_g1(self): | |
| """G1: Evaluate whether Area 1 should activate.""" | |
| stats1 = self.kg1.get_statistics() | |
| stats2 = self.kg2.get_statistics() | |
| return abs(stats1["avg_degree"] - stats2["avg_degree"]) > 2.0 | |
| def _run_l1_high_conf_selection(self, candidates): | |
| """L1: Select high-confidence alignment pairs.""" | |
| self._log("[L1] Heterogeneous Entity Selection...") | |
| high_conf = [] | |
| for eid1, cands in candidates.items(): | |
| if not cands: | |
| continue | |
| best = cands[0] | |
| gap = best["score"] - cands[1]["score"] if len(cands) > 1 else 1.0 | |
| if best["score"] > 0.3 and gap > 0.01: | |
| high_conf.append({ | |
| "eid1": eid1, | |
| "eid2": best["eid"], | |
| "name1": self.kg1.entities[eid1]["name"], | |
| "name2": self.kg2.entities[best["eid"]]["name"], | |
| "score": best["score"], | |
| "confidence": "high", | |
| }) | |
| self._log(f" Selected {len(high_conf)} high-confidence pairs") | |
| return high_conf | |
| def _run_s4_large_scale_alignment(self, high_conf_pairs): | |
| """S4: Large-scale alignment using structural embeddings.""" | |
| self._log("[S4] Large-Scale Accurate Alignment...") | |
| emb_kg1 = self._build_structural_embeddings(self.kg1) | |
| emb_kg2 = self._build_structural_embeddings(self.kg2) | |
| eids_kg1 = sorted(emb_kg1.keys()) | |
| eids_kg2 = sorted(emb_kg2.keys()) | |
| mat1 = np.array([emb_kg1[eid] for eid in eids_kg1]) | |
| mat2 = np.array([emb_kg2[eid] for eid in eids_kg2]) | |
| sim = cosine_similarity(mat1, mat2) | |
| high_conf_set = {(p["eid1"], p["eid2"]) for p in high_conf_pairs} | |
| uncertain_pairs = [] | |
| for i, eid1 in enumerate(eids_kg1): | |
| best_j = np.argmax(sim[i]) | |
| best_score = sim[i][best_j] | |
| eid2 = eids_kg2[best_j] | |
| if (eid1, eid2) not in high_conf_set and best_score > 0.2: | |
| uncertain_pairs.append({ | |
| "eid1": eid1, "eid2": eid2, | |
| "name1": self.kg1.entities[eid1]["name"], | |
| "name2": self.kg2.entities[eid2]["name"], | |
| "score": float(best_score), "confidence": "uncertain", | |
| }) | |
| self._log(f" Similarity matrix: {sim.shape}, Uncertain: {len(uncertain_pairs)}") | |
| return sim, uncertain_pairs | |
| def _build_structural_embeddings(self, kg, dim=64): | |
| """Build structural embeddings from degree and type features.""" | |
| embeddings = {} | |
| for eid in kg.entities: | |
| total_deg = kg.graph.degree(eid) | |
| in_deg = kg.graph.in_degree(eid) | |
| out_deg = kg.graph.out_degree(eid) | |
| vec = np.zeros(dim) | |
| vec[0] = np.log1p(total_deg) / 10.0 | |
| vec[1] = in_deg / max(total_deg, 1) | |
| vec[2] = out_deg / max(total_deg, 1) | |
| etype = kg.entities[eid].get("type", "Unknown") | |
| type_hash = hash(etype) % (dim - 3) | |
| vec[3 + type_hash] = 1.0 | |
| embeddings[eid] = vec | |
| return embeddings | |
| def _run_l2_meta_rules(self, high_conf_pairs, uncertain_pairs): | |
| """L2: Generate meta-logic symbol rules.""" | |
| self._log("[L2] Multi-Granularity Summarization...") | |
| rules = [ | |
| { | |
| "rule": "∃name(e1, n) ∧ ∃name(e2, n') ∧ similar(n, n') ⇒ align(e1, e2)", | |
| "description": "Entities with similar names are likely aligned", | |
| "confidence": 0.95, | |
| }, | |
| { | |
| "rule": "type(e1) = type(e2) ⇒ align(e1, e2)", | |
| "description": "Aligned entities tend to have the same type", | |
| "confidence": 0.88, | |
| }, | |
| { | |
| "rule": "struct_sim(e1, e2) > θ ⇒ align(e1, e2)", | |
| "description": "Entities with similar neighborhood structures are likely aligned", | |
| "confidence": 0.82, | |
| }, | |
| { | |
| "rule": "degree(e1) ≈ degree(e2) ⇒ align(e1, e2)", | |
| "description": "Aligned entities tend to have similar graph degrees", | |
| "confidence": 0.78, | |
| }, | |
| ] | |
| self._log(f" Generated {len(rules)} meta-logic rules") | |
| return rules | |
| def _run_l3_l4_knowledge_aggregation(self, uncertain_pairs, meta_rules): | |
| """L3/L4: Knowledge aggregation.""" | |
| self._log("[L3] Master — coordinating knowledge aggregation...") | |
| type_groups = defaultdict(list) | |
| for pair in uncertain_pairs: | |
| etype = self.kg1.entities.get(pair["eid1"], {}).get("type", "Unknown") | |
| type_groups[etype].append(pair) | |
| self._log(f" Assigned {len(type_groups)} domain expert groups") | |
| for etype, pairs in list(type_groups.items())[:5]: | |
| self._log(f" Domain Expert-{etype}: {len(pairs)} pairs") |