""" AgentEA-Studio: Real KG data generator. Generates realistic synthetic KG data for demo purposes, matching the statistics of DBP15K, ICEWS-WIKI, etc. """ import random import json import os from collections import defaultdict # Entity name pools for generating realistic KGs PERSON_NAMES = [ "Barack Obama", "Joe Biden", "Donald Trump", "Angela Merkel", "Emmanuel Macron", "Xi Jinping", "Vladimir Putin", "Boris Johnson", "Justin Trudeau", "Narendra Modi", "Mahmoud Abbas", "Benjamin Netanyahu", "Emmanuel Macron", "Ursula von der Leyen", "Antonio Guterres", "Steffi Lemke", "Olaf Scholz", "Yoshihide Suga", "Moon Jae-in", "Scott Morrison", "Jacinda Ardern", "Pedro Sanchez", "Mario Draghi", "Andrzej Duda", "Recep Erdogan", "Jair Bolsonaro", "Cyril Ramaphosa", "Alberto Fernandez", "Mette Frederiksen", "Sanna Marin", ] COUNTRY_NAMES = [ "United States", "China", "Russia", "Germany", "France", "United Kingdom", "Japan", "India", "Brazil", "Canada", "Australia", "Italy", "Spain", "South Korea", "Mexico", "Indonesia", "Turkey", "Netherlands", "Switzerland", "Sweden", "Poland", "Belgium", "Norway", "Austria", "Denmark", "Finland", "Portugal", "Greece", "Ireland", "New Zealand", "Singapore", "Malaysia", "Thailand", "Vietnam", "Philippines", "Egypt", "South Africa", "Nigeria", "Kenya", "Argentina", "Chile", "Colombia", "Peru", "Saudi Arabia", "UAE", "Qatar", "Israel", "Iran", "Pakistan", "Bangladesh", ] ORG_NAMES = [ "United Nations", "NATO", "European Union", "World Bank", "IMF", "World Health Organization", "UNESCO", "OPEC", "ASEAN", "African Union", "Red Cross", "Greenpeace", "Amnesty International", "World Trade Organization", "International Criminal Court", "OECD", "G20", "G7", "BRICS", "World Economic Forum", "International Olympic Committee", "FIFA", "Apple Inc.", "Google", "Microsoft", "Amazon", "Meta", "Tesla", "Samsung", "Toyota", "Volkswagen", "Shell", "ExxonMobil", "Pfizer", "Johnson & Johnson", "Roche", "Novartis", "Siemens", "IBM", "Intel", "NVIDIA", "AMD", "Qualcomm", "Huawei", "Alibaba", "Tencent", "Baidu", ] CITY_NAMES = [ "New York", "London", "Paris", "Tokyo", "Beijing", "Shanghai", "Moscow", "Berlin", "Washington D.C.", "Brussels", "Geneva", "Vienna", "Rome", "Madrid", "Lisbon", "Seoul", "Sydney", "Toronto", "Mumbai", "Delhi", "Dubai", "Singapore", "Hong Kong", "San Francisco", "Chicago", "Boston", "Los Angeles", "Houston", "Miami", "Seattle", "Austin", "Denver", "Cairo", "Lagos", "Nairobi", "Cape Town", "Buenos Aires", "Sao Paulo", "Mexico City", "Lima", "Bogota", "Santiago", "Bangkok", "Jakarta", "Manila", "Hanoi", "Kuala Lumpur", "Taipei", "Osaka", "Kyoto", "Zurich", "Stockholm", "Oslo", "Copenhagen", "Helsinki", "Warsaw", "Prague", "Budapest", "Athens", "Dublin", "Amsterdam", "Barcelona", "Munich", "Frankfurt", "Hamburg", "Milan", "Naples", "Venice", "Florence", ] EVENT_NAMES = [ "World War II", "Cold War", "Arab Spring", "Brexit", "COVID-19 Pandemic", "2008 Financial Crisis", "Paris Agreement", "Kyoto Protocol", "Iraq War", "Syrian Civil War", "Russo-Ukrainian War", "Eurozone Crisis", "Belt and Road Initiative", "Trans-Pacific Partnership", "NAFTA", "Space Race", "Moon Landing", "Chernobyl Disaster", "Fukushima Disaster", "Hurricane Katrina", "Indian Ocean Tsunami", "Haiti Earthquake", ] # Relation types RELATIONS = [ "president_of", "prime_minister_of", "located_in", "capital_of", "member_of", "part_of", "founded_by", "headquartered_in", "has_leader", "borders_with", "trades_with", "allied_with", "sanctions_against", "negotiates_with", "attended_by", "born_in", "died_in", "educated_at", "works_for", "owns", "subsidiary_of", "competitor_of", "supplier_to", "investor_in", "preceded_by", "followed_by", "caused_by", "resulted_in", "announced_at", "signed_at", " ratified_by", "participated_in", "voted_for", "vetoed_by", "proposed_by", "implemented_by", "reported_by", "published_in", "broadcast_by", "regulated_by", "funded_by", "audited_by", "certified_by", "accredited_by", ] # Entity types ENTITY_TYPES = ["Person", "Country", "City", "Organization", "Event", "Other"] # Word pools for generating diverse entity names in cross-lingual KGs FR_WORDS = ["président", "république", "organisation", "traité", "accord", "gouvernement", "ministre", "ambassadeur", "conseil", "commission"] ZH_TRANSLATIONS = { "United States": "美国", "China": "中国", "Russia": "俄罗斯", "Germany": "德国", "France": "法国", "United Kingdom": "英国", "Japan": "日本", "India": "印度", "Brazil": "巴西", "Canada": "加拿大", } class KGGenerator: """Generate realistic synthetic KG data for entity alignment demo.""" def __init__(self, seed=42): random.seed(seed) self.entity_counter = 0 self.rel_counter = 0 def _make_entity(self, name, etype, kg_prefix=""): eid = f"{kg_prefix}E{self.entity_counter}" self.entity_counter += 1 return {"id": eid, "name": name, "type": etype} def _make_relation(self, name): rid = f"R{self.rel_counter}" self.rel_counter += 1 return {"id": rid, "name": name} def generate_db15k_style(self, n_entities=5000, n_shared=5000): """Generate a DBP15K-style cross-lingual KG pair (EN-FR style). Two KGs with same entities but different languages/naming. """ entities_kg1 = [] entities_kg2 = [] relations = self._build_relation_set() triples_kg1 = [] triples_kg2 = [] alignments = [] # Create shared entities name_pools = [PERSON_NAMES, COUNTRY_NAMES, ORG_NAMES, CITY_NAMES] type_map = ["Person", "Country", "Organization", "City"] for i in range(n_shared): pool_idx = i % len(name_pools) name = random.choice(name_pools[pool_idx]) etype = type_map[pool_idx] e1 = self._make_entity(name, etype, "EN_") e2 = self._make_entity(self._frenchify(name), etype, "FR_") entities_kg1.append(e1) entities_kg2.append(e2) alignments.append((e1["id"], e2["id"])) # Add KG-specific entities for i in range(n_entities - n_shared): etype = random.choice(type_map) pool = name_pools[type_map.index(etype)] name = random.choice(pool) e1 = self._make_entity(name, etype, "EN_") entities_kg1.append(e1) for i in range(n_entities - n_shared): etype = random.choice(type_map) pool = name_pools[type_map.index(etype)] name = random.choice(pool) e2 = self._make_entity(self._frenchify(name), etype, "FR_") entities_kg2.append(e2) # Generate triples triples_kg1 = self._generate_triples(entities_kg1, relations, avg_degree=6) triples_kg2 = self._generate_triples(entities_kg2, relations, avg_degree=6) return { "kg1": {"entities": entities_kg1, "relations": relations, "triples": triples_kg1}, "kg2": {"entities": entities_kg2, "relations": relations, "triples": triples_kg2}, "alignments": alignments, } def generate_hhkg_style(self, n_entities_kg1=3000, n_entities_kg2=5000, n_shared=1500): """Generate an HHKG-style KG pair (ICEWS-WIKI style). Two KGs with very different scales, densities, and structures. KG1 is dense (event-focused), KG2 is sparse (encyclopedic). """ entities_kg1 = [] entities_kg2 = [] relations = self._build_relation_set() triples_kg1 = [] triples_kg2 = [] alignments = [] name_pools = [PERSON_NAMES, COUNTRY_NAMES, ORG_NAMES, CITY_NAMES, EVENT_NAMES] type_map = ["Person", "Country", "Organization", "City", "Event"] # Shared entities for i in range(n_shared): pool_idx = i % min(len(name_pools), 4) # Keep shared entities to common types name = random.choice(name_pools[pool_idx]) etype = type_map[pool_idx] e1 = self._make_entity(name, etype, "KG1_") e2 = self._make_entity(f"{name} (wiki)", etype, "KG2_") entities_kg1.append(e1) entities_kg2.append(e2) alignments.append((e1["id"], e2["id"])) # KG1-specific entities (dense, event-focused) for i in range(n_entities_kg1 - n_shared): etype = random.choice(["Event", "Person", "Country"]) pool = name_pools[type_map.index(etype)] name = random.choice(pool) e1 = self._make_entity(name, etype, "KG1_") entities_kg1.append(e1) # KG2-specific entities (sparse, encyclopedic) for i in range(n_entities_kg2 - n_shared): etype = random.choice(type_map) pool = name_pools[type_map.index(etype)] name = random.choice(pool) e2 = self._make_entity(f"{name} (wiki)", etype, "KG2_") entities_kg2.append(e2) # KG1: very dense (many triples per entity) triples_kg1 = self._generate_triples(entities_kg1, relations, avg_degree=25) # KG2: sparse (few triples per entity) triples_kg2 = self._generate_triples(entities_kg2, relations, avg_degree=4) return { "kg1": {"entities": entities_kg1, "relations": relations, "triples": triples_kg1}, "kg2": {"entities": entities_kg2, "relations": relations, "triples": triples_kg2}, "alignments": alignments, } def _build_relation_set(self): """Build a relation vocabulary.""" relations = [] for rel_name in RELATIONS: relations.append(self._make_relation(rel_name)) return relations def _generate_triples(self, entities, relations, avg_degree=6): """Generate realistic triples connecting entities.""" triples = [] n_entities = len(entities) total_triples = int(n_entities * avg_degree) # Build relation index rel_by_name = {r["name"]: r for r in relations} # Assign types entities_by_type = defaultdict(list) for e in entities: entities_by_type[e["type"]].append(e) for _ in range(total_triples): head = random.choice(entities) # Choose relation appropriate for the head entity type etype = head["type"] if etype == "Person": possible_rels = ["president_of", "prime_minister_of", "born_in", "educated_at", "works_for", "has_leader", "attended_by", "participated_in"] tail_pool = entities_by_type["Country"] + entities_by_type["City"] + \ entities_by_type["Organization"] + entities_by_type["Event"] elif etype == "Country": possible_rels = ["located_in", "capital_of", "member_of", "borders_with", "trades_with", "allied_with", "sanctions_against", "negotiates_with", "part_of", "has_leader"] tail_pool = entities_by_type["Country"] + entities_by_type["City"] + \ entities_by_type["Organization"] + entities_by_type["Person"] elif etype == "City": possible_rels = ["located_in", "capital_of", "part_of", "headquartered_in", "born_in", "died_in"] tail_pool = entities_by_type["Country"] + entities_by_type["Organization"] + \ entities_by_type["Person"] elif etype == "Organization": possible_rels = ["headquartered_in", "member_of", "part_of", "founded_by", "subsidiary_of", "competitor_of", "funded_by", "regulated_by", "participated_in"] tail_pool = entities_by_type["City"] + entities_by_type["Country"] + \ entities_by_type["Organization"] + entities_by_type["Person"] elif etype == "Event": possible_rels = ["attended_by", "participated_in", "caused_by", "resulted_in", "announced_at", "signed_at", "reported_by", "preceded_by", "followed_by"] tail_pool = entities_by_type["Person"] + entities_by_type["Country"] + \ entities_by_type["Organization"] + entities_by_type["City"] + \ entities_by_type["Event"] else: possible_rels = list(RELATIONS) tail_pool = entities if not tail_pool: continue rel_name = random.choice([r for r in possible_rels if r in rel_by_name]) tail = random.choice(tail_pool) if tail["id"] != head["id"]: triples.append({ "head": head["id"], "relation": rel_by_name[rel_name]["id"], "tail": tail["id"], }) return triples def _frenchify(self, name): """Create a French-like version of an English entity name.""" if name in ZH_TRANSLATIONS: return name parts = name.split() if len(parts) > 1: return f"{parts[-1]}, {' '.join(parts[:-1])}" return name def save_dataset(self, dataset, output_dir): """Save a generated dataset to disk.""" os.makedirs(output_dir, exist_ok=True) for kg_name in ["kg1", "kg2"]: kg = dataset[kg_name] kg_dir = os.path.join(output_dir, kg_name) os.makedirs(kg_dir, exist_ok=True) # Save entities with open(os.path.join(kg_dir, "entities.json"), "w") as f: json.dump(kg["entities"], f, indent=2) # Save relations with open(os.path.join(kg_dir, "relations.json"), "w") as f: json.dump(kg["relations"], f, indent=2) # Save triples with open(os.path.join(kg_dir, "triples.json"), "w") as f: json.dump(kg["triples"], f, indent=2) # Save triples as TSV for easy loading with open(os.path.join(kg_dir, "triples.tsv"), "w") as f: for t in kg["triples"]: f.write(f"{t['head']}\t{t['relation']}\t{t['tail']}\n") # Save alignments with open(os.path.join(output_dir, "alignments.json"), "w") as f: json.dump(dataset["alignments"], f, indent=2) def generate_all_datasets(output_dir): """Generate all demo datasets.""" gen = KGGenerator(seed=42) print("Generating DBP15K-style dataset...") db15k = gen.generate_db15k_style(n_entities=3000, n_shared=3000) gen.save_dataset(db15k, os.path.join(output_dir, "db15k_demo")) print("Generating HHKG-style dataset (ICEWS-WIKI style)...") gen2 = KGGenerator(seed=123) hhkg = gen2.generate_hhkg_style(n_entities_kg1=2000, n_entities_kg2=3500, n_shared=1000) gen2.save_dataset(hhkg, os.path.join(output_dir, "hhkg_demo")) print("Done!") # Print statistics for name, ds in [("DBP15K-Demo", db15k), ("HHKG-Demo", hhkg)]: kg1 = ds["kg1"] kg2 = ds["kg2"] print(f"\n{name}:") print(f" KG1: {len(kg1['entities'])} entities, {len(kg1['triples'])} triples, " f"density={len(kg1['triples'])/len(kg1['entities']):.1f}") print(f" KG2: {len(kg2['entities'])} entities, {len(kg2['triples'])} triples, " f"density={len(kg2['triples'])/len(kg2['entities']):.1f}") print(f" Alignments: {len(ds['alignments'])}") if __name__ == "__main__": output_dir = os.path.join(os.path.dirname(__file__), "..", "data") generate_all_datasets(output_dir)