""" 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")