EvoRM / experiments /run_all_experiments.py

Security: remove API key from experiments/run_all_experiments.py

cae54b6 verified 3 months ago

22.1 kB

#!/usr/bin/env python3
"""
EvoRM Comprehensive Experiment Runner
=======================================
Covers all 5 Research Questions from the TKDE paper:
  RQ1: Main Results (different datasets)
  RQ2: Cross-Model Analysis
  RQ3: Ablation Study
  RQ4: Efficiency Analysis
  RQ5: Cold-Start Scaling
Usage:
  python run_all_experiments.py --all          # Run everything
  python run_all_experiments.py --ablation     # RQ3 only
  python run_all_experiments.py --coldstart    # RQ5 only
  python run_all_experiments.py --main         # RQ1 only
  python run_all_experiments.py --efficiency   # RQ4 only
"""

import os
import sys
import json
import time
import shutil
import argparse
import subprocess
from datetime import datetime
from collections import defaultdict

# Add parent dir to path
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

os.environ["OPENAI_API_BASE"] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
# OPENAI_API_KEY set via environment variable

RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
DATA_BASE = "/root/autodl-tmp/AdaCoAgentEA/data"

# Dataset configurations
DATASETS = {
    'icews_wiki': {
        'path': f'{DATA_BASE}/icews_wiki',
        'ref_pairs': f'{DATA_BASE}/icews_wiki/ref_pairs',
        'n_pairs': 3540,
        'type': 'HHEA',
    },
    # ICEWS-YAGO and BETA require full pipeline (Area1→Area2→Area3)
    # They only have index/role_assignments, not raw KG triples
}

def ensure_dir(d):
    os.makedirs(d, exist_ok=True)

def load_ref_pairs(path):
    """Load reference entity pairs."""
    pairs = []
    with open(path, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                pairs.append((int(parts[0]), int(parts[1])))
    return pairs

def compute_metrics(aligned_file, ref_pairs_file):
    """Compute Hits@1, MRR, Precision, Recall."""
    # Load aligned pairs
    aligned = set()
    if os.path.exists(aligned_file):
        with open(aligned_file, 'r') as f:
            for line in f:
                parts = line.strip().split('\t')
                if len(parts) == 2:
                    try:
                        aligned.add((int(parts[0]), int(parts[1])))
                    except ValueError:
                        continue

    # Load reference pairs
    ref = set()
    with open(ref_pairs_file, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    ref.add((int(parts[0]), int(parts[1])))
                except ValueError:
                    continue

    if not ref:
        return {'error': 'No reference pairs found'}

    # Compute metrics
    true_positives = aligned & ref
    false_positives = aligned - ref
    false_negatives = ref - aligned

    tp = len(true_positives)
    fp = len(false_positives)
    fn = len(false_negatives)

    precision = tp / max(1, tp + fp)
    recall = tp / max(1, tp + fn)
    hits1 = tp / len(ref)  # Hits@1 since we have 1-1 alignment
    mrr = hits1  # MRR = Hits@1 for 1-1 alignment

    return {
        'Hits@1': round(hits1, 4),
        'MRR': round(mrr, 4),
        'Precision': round(precision, 4),
        'Recall': round(recall, 4),
        'TP': tp, 'FP': fp, 'FN': fn,
        'Total_Ref': len(ref),
        'Total_Aligned': len(aligned),
    }

def run_llm1_alignment(data_dir, from_m3=False, ablation_config=None, 
                        no_optimization_tool=False, n_entities=None):
    """Run LLM1 label selector alignment."""
    from LLM1_label_selector import align_entities
    
    # Backup original important_entities if we're limiting entities
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    
    if n_entities:
        # Create subset
        bak = important_file + ".bak_full"
        if not os.path.exists(bak):
            shutil.copy(important_file, bak)
        # Read first n_entities
        with open(bak, 'r') as f:
            lines = f.readlines()
        with open(important_file, 'w') as f:
            f.writelines(lines[:n_entities])
    
    try:
        start_time = time.time()
        aligned_pairs = align_entities(
            data_dir, 
            from_m3=from_m3,
            ablation_config=ablation_config,
            no_optimization_tool=no_optimization_tool
        )
        elapsed = time.time() - start_time
        
        # Get token count
        import tokens_cal
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0  # Reset
        
        return {
            'aligned_pairs': len(aligned_pairs),
            'time_seconds': elapsed,
            'tokens': tokens,
        }
    finally:
        # Restore
        if n_entities and os.path.exists(bak):
            shutil.copy(bak, important_file)

def run_ablation_study():
    """RQ3: Ablation Study - Full System vs w/o Stage-1 vs w/o Maintenance."""
    print("\n" + "="*70)
    print("RQ3: ABLATION STUDY")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    n_entities = 200  # Subset for efficiency
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Set up important_entities subset
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    # Read reference pairs to get entity IDs
    ref_set = load_ref_pairs(ref_pairs)
    kg1_entities = list(set(p[0] for p in ref_set))
    
    # Create important_entities with first n_entities
    with open(important_file, 'w') as f:
        for eid in kg1_entities[:n_entities]:
            f.write(f"{eid}\n")
    
    # Also update aligned_entities to only contain relevant ref pairs
    with open(important_file, 'r') as f:
        selected_entities = set(int(line.strip()) for line in f if line.strip())
    
    results = {}
    
    configs = [
        ('Full System (EvoRM)', None),
        ('w/o Stage-1 (Symbolic)', 'no_stage1'),
        ('w/o Maintenance', 'no_maintenance'),
    ]
    
    for config_name, ablation_mode in configs:
        print(f"\n{'='*50}")
        print(f"Config: {config_name}")
        print(f"{'='*50}")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        # Clear aligned_entities
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Reset tokens
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        # Run alignment with ablation config
        # We need to patch the module to accept ablation_mode
        import evorm_plugin
        
        # Monkey-patch: create a new EvoRMPlugin with ablation_mode
        original_init = evorm_plugin.EvoRMPlugin.__init__
        
        # Run using LLM1_label_selector directly
        from LLM1_label_selector import align_entities
        
        # We need to inject ablation_mode into the EvoRMPlugin initialization
        # The simplest way is to modify the module temporarily
        import LLM1_label_selector as llm1_module
        
        # Save original
        orig_align = llm1_module.align_entities
        
        # Create a wrapper
        def make_align_with_ablation(ablation):
            def wrapped_align(data_dir, from_m3=False, ablation_config=None, no_optimization_tool=False):
                # Override evorm_plugin to use ablation
                original_plugin = evorm_plugin.EvoRMPlugin
                class AblationPlugin(original_plugin):
                    def __init__(self, **kwargs):
                        kwargs['ablation_mode'] = ablation
                        super().__init__(**kwargs)
                evorm_plugin.EvoRMPlugin = AblationPlugin
                try:
                    return orig_align(data_dir, from_m3, ablation_config, no_optimization_tool)
                finally:
                    evorm_plugin.EvoRMPlugin = original_plugin
            return wrapped_align
        
        if ablation_mode:
            llm1_module.align_entities = make_align_with_ablation(ablation_mode)
        
        try:
            start_time = time.time()
            aligned_pairs = llm1_module.align_entities(data_dir, from_m3=False)
            elapsed = time.time() - start_time
            
            tokens = tokens_cal.global_tokens
            tokens_cal.global_tokens = 0
            
            # Compute metrics
            metrics = compute_metrics(aligned_file, ref_pairs)
            
            results[config_name] = {
                'aligned_pairs': len(aligned_pairs),
                'time_seconds': elapsed,
                'tokens': tokens,
                'metrics': metrics,
            }
            
            print(f"  Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s, Tokens: {tokens}")
            print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}, MRR: {metrics.get('MRR', 'N/A')}")
            print(f"  Precision: {metrics.get('Precision', 'N/A')}, Recall: {metrics.get('Recall', 'N/A')}")
            
        finally:
            if ablation_mode:
                llm1_module.align_entities = orig_align
    
    # Restore full important_entities
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"ablation_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nAblation results saved to: {result_file}")
    
    # Print summary table
    print("\n" + "="*70)
    print("ABLATION STUDY SUMMARY")
    print("="*70)
    print(f"{'Configuration':<30} {'Hits@1':<10} {'MRR':<10} {'Precision':<10} {'Recall':<10} {'Tokens':<10}")
    print("-"*70)
    for config_name in ['Full System (EvoRM)', 'w/o Stage-1 (Symbolic)', 'w/o Maintenance']:
        if config_name in results:
            r = results[config_name]
            m = r['metrics']
            print(f"{config_name:<30} {m.get('Hits@1', 0):<10.4f} {m.get('MRR', 0):<10.4f} {m.get('Precision', 0):<10.4f} {m.get('Recall', 0):<10.4f} {r['tokens']:<10}")
    
    return results

def run_coldstart_experiment():
    """RQ5: Cold-Start Scaling - different warmup sizes."""
    print("\n" + "="*70)
    print("RQ5: COLD-START SCALING EXPERIMENT")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Read all reference pairs for ground truth
    ref_set = load_ref_pairs(ref_pairs)
    kg1_entities = list(set(p[0] for p in ref_set))
    
    # Important entities file
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    results = {}
    
    warmup_sizes = [100, 200, 500, 1000, 2000]
    
    for n_warmup in warmup_sizes:
        print(f"\n--- N_warmup = {n_warmup} ---")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        # Clear aligned_entities
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Create subset
        with open(important_file, 'w') as f:
            for eid in kg1_entities[:n_warmup]:
                f.write(f"{eid}\n")
        
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        from LLM1_label_selector import align_entities
        
        start_time = time.time()
        aligned_pairs = align_entities(data_dir, from_m3=False)
        elapsed = time.time() - start_time
        
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0
        
        metrics = compute_metrics(aligned_file, ref_pairs)
        
        results[f'n_warmup_{n_warmup}'] = {
            'aligned_pairs': len(aligned_pairs),
            'time_seconds': elapsed,
            'tokens': tokens,
            'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
            'metrics': metrics,
        }
        
        print(f"  Aligned: {len(aligned_pairs)}, Tokens: {tokens}, "
              f"Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
        print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}")
    
    # Restore
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save results
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"coldstart_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nCold-start results saved to: {result_file}")
    
    # Print summary
    print("\n" + "="*70)
    print("COLD-START SCALING SUMMARY")
    print("="*70)
    print(f"{'N_warmup':<12} {'Hits@1':<10} {'Tokens':<12} {'Tokens/Pair':<12} {'Time(s)':<10}")
    print("-"*56)
    for n_warmup in warmup_sizes:
        key = f'n_warmup_{n_warmup}'
        if key in results:
            r = results[key]
            m = r['metrics']
            print(f"{n_warmup:<12} {m.get('Hits@1', 0):<10.4f} {r['tokens']:<12} {r['tokens_per_pair']:<12.1f} {r['time_seconds']:<10.1f}")
    
    return results

def run_efficiency_analysis():
    """RQ4: Efficiency Analysis - token consumption and latency."""
    print("\n" + "="*70)
    print("RQ4: EFFICIENCY ANALYSIS")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
    bak = important_file + ".bak_full"
    if not os.path.exists(bak):
        shutil.copy(important_file, bak)
    
    results = {}
    
    # Test with different entity counts
    entity_counts = [100, 200, 500]
    
    for n_entities in entity_counts:
        print(f"\n--- N_entities = {n_entities} ---")
        
        # Clear evorm state
        evorm_state = os.path.join(data_dir, "evorm_state")
        if os.path.exists(evorm_state):
            shutil.rmtree(evorm_state)
        
        if os.path.exists(aligned_file):
            os.remove(aligned_file)
        
        # Subset
        ref_set = load_ref_pairs(ref_pairs)
        kg1_entities = list(set(p[0] for p in ref_set))
        with open(important_file, 'w') as f:
            for eid in kg1_entities[:n_entities]:
                f.write(f"{eid}\n")
        
        import tokens_cal
        tokens_cal.global_tokens = 0
        
        from LLM1_label_selector import align_entities
        
        # Measure latency
        latencies = []
        original_task = None
        
        start_time = time.time()
        aligned_pairs = align_entities(data_dir, from_m3=False)
        total_time = time.time() - start_time
        
        tokens = tokens_cal.global_tokens
        tokens_cal.global_tokens = 0
        
        metrics = compute_metrics(aligned_file, ref_pairs)
        
        n_pairs = len(aligned_pairs)
        results[f'n_{n_entities}'] = {
            'n_pairs': n_pairs,
            'total_time_s': total_time,
            'time_per_pair_s': total_time / max(1, n_pairs),
            'tokens': tokens,
            'tokens_per_pair': tokens / max(1, n_pairs),
            'metrics': metrics,
        }
        
        print(f"  Pairs: {n_pairs}, Time: {total_time:.1f}s, "
              f"Time/Pair: {total_time/max(1,n_pairs):.2f}s, "
              f"Tokens/Pair: {tokens/max(1,n_pairs):.1f}")
    
    # Restore
    if os.path.exists(bak):
        shutil.copy(bak, important_file)
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"efficiency_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"\nEfficiency results saved to: {result_file}")
    
    # Summary
    print("\n" + "="*70)
    print("EFFICIENCY SUMMARY")
    print("="*70)
    print(f"{'N_Entities':<12} {'Pairs':<8} {'Time(s)':<10} {'Time/Pair':<12} {'Tokens/Pair':<12} {'Hits@1':<10}")
    print("-"*64)
    for n_entities in entity_counts:
        key = f'n_{n_entities}'
        if key in results:
            r = results[key]
            m = r['metrics']
            print(f"{n_entities:<12} {r['n_pairs']:<8} {r['total_time_s']:<10.1f} {r['time_per_pair_s']:<12.2f} {r['tokens_per_pair']:<12.1f} {m.get('Hits@1', 0):<10.4f}")
    
    return results

def run_main_experiment():
    """RQ1: Main Results - full dataset run."""
    print("\n" + "="*70)
    print("RQ1: MAIN EXPERIMENT - ICEWS-WIKI (Full Dataset)")
    print("="*70)
    
    data_dir = f"{DATA_BASE}/icews_wiki"
    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
    ref_pairs = os.path.join(data_dir, "ref_pairs")
    
    # Clear evorm state
    evorm_state = os.path.join(data_dir, "evorm_state")
    if os.path.exists(evorm_state):
        shutil.rmtree(evorm_state)
    
    if os.path.exists(aligned_file):
        os.remove(aligned_file)
    
    import tokens_cal
    tokens_cal.global_tokens = 0
    
    from LLM1_label_selector import align_entities
    
    start_time = time.time()
    aligned_pairs = align_entities(data_dir, from_m3=False)
    elapsed = time.time() - start_time
    
    tokens = tokens_cal.global_tokens
    tokens_cal.global_tokens = 0
    
    metrics = compute_metrics(aligned_file, ref_pairs)
    
    results = {
        'dataset': 'ICEWS-WIKI',
        'n_ref_pairs': 3540,
        'aligned_pairs': len(aligned_pairs),
        'time_seconds': elapsed,
        'tokens': tokens,
        'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
        'metrics': metrics,
    }
    
    # Check evorm state
    if os.path.exists(evorm_state):
        rules_file = os.path.join(evorm_state, "rules.json")
        if os.path.exists(rules_file):
            with open(rules_file) as f:
                rules = json.load(f)
            results['n_rules'] = len(rules)
        hyperedges_file = os.path.join(evorm_state, "hyperedges.json")
        if os.path.exists(hyperedges_file):
            with open(hyperedges_file) as f:
                hedges = json.load(f)
            results['n_hyperedges'] = len(hedges)
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"main_icews_wiki_{timestamp}.json")
    ensure_dir(RESULTS_DIR)
    with open(result_file, 'w') as f:
        json.dump(results, f, indent=2)
    
    print(f"Main results saved to: {result_file}")
    print(f"\n  Dataset: ICEWS-WIKI")
    print(f"  Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s")
    print(f"  Tokens: {tokens}, Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
    print(f"  Hits@1: {metrics.get('Hits@1', 'N/A')}")
    print(f"  MRR: {metrics.get('MRR', 'N/A')}")
    print(f"  Precision: {metrics.get('Precision', 'N/A')}")
    print(f"  Recall: {metrics.get('Recall', 'N/A')}")
    if 'n_rules' in results:
        print(f"  Rules: {results['n_rules']}, Hyperedges: {results.get('n_hyperedges', 'N/A')}")
    
    return results

def generate_latex_table(all_results):
    """Generate paper-style LaTeX table."""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    table_file = os.path.join(RESULTS_DIR, f"paper_table_{timestamp}.tex")
    
    latex = r"""\begin{table}[t]
\centering
\caption{Experimental Results of EvoRM on ICEWS-WIKI}
\label{tab:main_results}
\begin{tabular}{lcccc}
\hline
\textbf{Method} & \textbf{Hits@1} & \textbf{MRR} & \textbf{Precision} & \textbf{Recall} \\
\hline
"""
    
    # Main results
    if 'main' in all_results:
        m = all_results['main']['metrics']
        latex += f"EvoRM (Ours) & {m['Hits@1']:.4f} & {m['MRR']:.4f} & {m['Precision']:.4f} & {m['Recall']:.4f} \\\\\n"
    
    latex += r"""\hline
\end{tabular}
\end{table}
"""
    
    with open(table_file, 'w') as f:
        f.write(latex)
    
    print(f"LaTeX table saved to: {table_file}")
    return table_file

def main():
    parser = argparse.ArgumentParser(description='EvoRM Comprehensive Experiment Runner')
    parser.add_argument('--all', action='store_true', help='Run all experiments')
    parser.add_argument('--ablation', action='store_true', help='Run ablation study (RQ3)')
    parser.add_argument('--coldstart', action='store_true', help='Run cold-start scaling (RQ5)')
    parser.add_argument('--main', action='store_true', help='Run main experiment (RQ1)')
    parser.add_argument('--efficiency', action='store_true', help='Run efficiency analysis (RQ4)')
    args = parser.parse_args()
    
    # Default: run everything if no flag specified
    run_all = args.all or not any([args.ablation, args.coldstart, args.main, args.efficiency])
    
    all_results = {}
    
    if run_all or args.main:
        all_results['main'] = run_main_experiment()
    
    if run_all or args.ablation:
        all_results['ablation'] = run_ablation_study()
    
    if run_all or args.coldstart:
        all_results['coldstart'] = run_coldstart_experiment()
    
    if run_all or args.efficiency:
        all_results['efficiency'] = run_efficiency_analysis()
    
    # Generate paper table
    generate_latex_table(all_results)
    
    print("\n" + "="*70)
    print("ALL EXPERIMENTS COMPLETE!")
    print("="*70)
    print(f"Results saved in: {RESULTS_DIR}")
    for f in sorted(os.listdir(RESULTS_DIR)):
        print(f"  {f}")

if __name__ == '__main__':
    main()

	#!/usr/bin/env python3
	"""
	EvoRM Comprehensive Experiment Runner
	=======================================
	Covers all 5 Research Questions from the TKDE paper:
	RQ1: Main Results (different datasets)
	RQ2: Cross-Model Analysis
	RQ3: Ablation Study
	RQ4: Efficiency Analysis
	RQ5: Cold-Start Scaling

	Usage:
	python run_all_experiments.py --all # Run everything
	python run_all_experiments.py --ablation # RQ3 only
	python run_all_experiments.py --coldstart # RQ5 only
	python run_all_experiments.py --main # RQ1 only
	python run_all_experiments.py --efficiency # RQ4 only
	"""

	import os
	import sys
	import json
	import time
	import shutil
	import argparse
	import subprocess
	from datetime import datetime
	from collections import defaultdict

	# Add parent dir to path
	sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
	sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')

	os.environ["OPENAI_API_BASE"] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
	# OPENAI_API_KEY set via environment variable

	RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
	DATA_BASE = "/root/autodl-tmp/AdaCoAgentEA/data"

	# Dataset configurations
	DATASETS = {
	'icews_wiki': {
	'path': f'{DATA_BASE}/icews_wiki',
	'ref_pairs': f'{DATA_BASE}/icews_wiki/ref_pairs',
	'n_pairs': 3540,
	'type': 'HHEA',
	},
	# ICEWS-YAGO and BETA require full pipeline (Area1→Area2→Area3)
	# They only have index/role_assignments, not raw KG triples
	}

	def ensure_dir(d):
	os.makedirs(d, exist_ok=True)

	def load_ref_pairs(path):
	"""Load reference entity pairs."""
	pairs = []
	with open(path, 'r') as f:
	for line in f:
	parts = line.strip().split('\t')
	if len(parts) == 2:
	pairs.append((int(parts[0]), int(parts[1])))
	return pairs

	def compute_metrics(aligned_file, ref_pairs_file):
	"""Compute Hits@1, MRR, Precision, Recall."""
	# Load aligned pairs
	aligned = set()
	if os.path.exists(aligned_file):
	with open(aligned_file, 'r') as f:
	for line in f:
	parts = line.strip().split('\t')
	if len(parts) == 2:
	try:
	aligned.add((int(parts[0]), int(parts[1])))
	except ValueError:
	continue

	# Load reference pairs
	ref = set()
	with open(ref_pairs_file, 'r') as f:
	for line in f:
	parts = line.strip().split('\t')
	if len(parts) == 2:
	try:
	ref.add((int(parts[0]), int(parts[1])))
	except ValueError:
	continue

	if not ref:
	return {'error': 'No reference pairs found'}

	# Compute metrics
	true_positives = aligned & ref
	false_positives = aligned - ref
	false_negatives = ref - aligned

	tp = len(true_positives)
	fp = len(false_positives)
	fn = len(false_negatives)

	precision = tp / max(1, tp + fp)
	recall = tp / max(1, tp + fn)
	hits1 = tp / len(ref) # Hits@1 since we have 1-1 alignment
	mrr = hits1 # MRR = Hits@1 for 1-1 alignment

	return {
	'Hits@1': round(hits1, 4),
	'MRR': round(mrr, 4),
	'Precision': round(precision, 4),
	'Recall': round(recall, 4),
	'TP': tp, 'FP': fp, 'FN': fn,
	'Total_Ref': len(ref),
	'Total_Aligned': len(aligned),
	}

	def run_llm1_alignment(data_dir, from_m3=False, ablation_config=None,
	no_optimization_tool=False, n_entities=None):
	"""Run LLM1 label selector alignment."""
	from LLM1_label_selector import align_entities

	# Backup original important_entities if we're limiting entities
	important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")

	if n_entities:
	# Create subset
	bak = important_file + ".bak_full"
	if not os.path.exists(bak):
	shutil.copy(important_file, bak)
	# Read first n_entities
	with open(bak, 'r') as f:
	lines = f.readlines()
	with open(important_file, 'w') as f:
	f.writelines(lines[:n_entities])

	try:
	start_time = time.time()
	aligned_pairs = align_entities(
	data_dir,
	from_m3=from_m3,
	ablation_config=ablation_config,
	no_optimization_tool=no_optimization_tool
	)
	elapsed = time.time() - start_time

	# Get token count
	import tokens_cal
	tokens = tokens_cal.global_tokens
	tokens_cal.global_tokens = 0 # Reset

	return {
	'aligned_pairs': len(aligned_pairs),
	'time_seconds': elapsed,
	'tokens': tokens,
	}
	finally:
	# Restore
	if n_entities and os.path.exists(bak):
	shutil.copy(bak, important_file)

	def run_ablation_study():
	"""RQ3: Ablation Study - Full System vs w/o Stage-1 vs w/o Maintenance."""
	print("\n" + "="*70)
	print("RQ3: ABLATION STUDY")
	print("="*70)

	data_dir = f"{DATA_BASE}/icews_wiki"
	n_entities = 200 # Subset for efficiency
	aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
	ref_pairs = os.path.join(data_dir, "ref_pairs")

	# Set up important_entities subset
	important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
	bak = important_file + ".bak_full"
	if not os.path.exists(bak):
	shutil.copy(important_file, bak)

	# Read reference pairs to get entity IDs
	ref_set = load_ref_pairs(ref_pairs)
	kg1_entities = list(set(p[0] for p in ref_set))

	# Create important_entities with first n_entities
	with open(important_file, 'w') as f:
	for eid in kg1_entities[:n_entities]:
	f.write(f"{eid}\n")

	# Also update aligned_entities to only contain relevant ref pairs
	with open(important_file, 'r') as f:
	selected_entities = set(int(line.strip()) for line in f if line.strip())

	results = {}

	configs = [
	('Full System (EvoRM)', None),
	('w/o Stage-1 (Symbolic)', 'no_stage1'),
	('w/o Maintenance', 'no_maintenance'),
	]

	for config_name, ablation_mode in configs:
	print(f"\n{'='*50}")
	print(f"Config: {config_name}")
	print(f"{'='*50}")

	# Clear evorm state
	evorm_state = os.path.join(data_dir, "evorm_state")
	if os.path.exists(evorm_state):
	shutil.rmtree(evorm_state)

	# Clear aligned_entities
	if os.path.exists(aligned_file):
	os.remove(aligned_file)

	# Reset tokens
	import tokens_cal
	tokens_cal.global_tokens = 0

	# Run alignment with ablation config
	# We need to patch the module to accept ablation_mode
	import evorm_plugin

	# Monkey-patch: create a new EvoRMPlugin with ablation_mode
	original_init = evorm_plugin.EvoRMPlugin.__init__

	# Run using LLM1_label_selector directly
	from LLM1_label_selector import align_entities

	# We need to inject ablation_mode into the EvoRMPlugin initialization
	# The simplest way is to modify the module temporarily
	import LLM1_label_selector as llm1_module

	# Save original
	orig_align = llm1_module.align_entities

	# Create a wrapper
	def make_align_with_ablation(ablation):
	def wrapped_align(data_dir, from_m3=False, ablation_config=None, no_optimization_tool=False):
	# Override evorm_plugin to use ablation
	original_plugin = evorm_plugin.EvoRMPlugin
	class AblationPlugin(original_plugin):
	def __init__(self, **kwargs):
	kwargs['ablation_mode'] = ablation
	super().__init__(**kwargs)
	evorm_plugin.EvoRMPlugin = AblationPlugin
	try:
	return orig_align(data_dir, from_m3, ablation_config, no_optimization_tool)
	finally:
	evorm_plugin.EvoRMPlugin = original_plugin
	return wrapped_align

	if ablation_mode:
	llm1_module.align_entities = make_align_with_ablation(ablation_mode)

	try:
	start_time = time.time()
	aligned_pairs = llm1_module.align_entities(data_dir, from_m3=False)
	elapsed = time.time() - start_time

	tokens = tokens_cal.global_tokens
	tokens_cal.global_tokens = 0

	# Compute metrics
	metrics = compute_metrics(aligned_file, ref_pairs)

	results[config_name] = {
	'aligned_pairs': len(aligned_pairs),
	'time_seconds': elapsed,
	'tokens': tokens,
	'metrics': metrics,
	}

	print(f" Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s, Tokens: {tokens}")
	print(f" Hits@1: {metrics.get('Hits@1', 'N/A')}, MRR: {metrics.get('MRR', 'N/A')}")
	print(f" Precision: {metrics.get('Precision', 'N/A')}, Recall: {metrics.get('Recall', 'N/A')}")

	finally:
	if ablation_mode:
	llm1_module.align_entities = orig_align

	# Restore full important_entities
	if os.path.exists(bak):
	shutil.copy(bak, important_file)

	# Save results
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	result_file = os.path.join(RESULTS_DIR, f"ablation_{timestamp}.json")
	ensure_dir(RESULTS_DIR)
	with open(result_file, 'w') as f:
	json.dump(results, f, indent=2)

	print(f"\nAblation results saved to: {result_file}")

	# Print summary table
	print("\n" + "="*70)
	print("ABLATION STUDY SUMMARY")
	print("="*70)
	print(f"{'Configuration':<30} {'Hits@1':<10} {'MRR':<10} {'Precision':<10} {'Recall':<10} {'Tokens':<10}")
	print("-"*70)
	for config_name in ['Full System (EvoRM)', 'w/o Stage-1 (Symbolic)', 'w/o Maintenance']:
	if config_name in results:
	r = results[config_name]
	m = r['metrics']
	print(f"{config_name:<30} {m.get('Hits@1', 0):<10.4f} {m.get('MRR', 0):<10.4f} {m.get('Precision', 0):<10.4f} {m.get('Recall', 0):<10.4f} {r['tokens']:<10}")

	return results

	def run_coldstart_experiment():
	"""RQ5: Cold-Start Scaling - different warmup sizes."""
	print("\n" + "="*70)
	print("RQ5: COLD-START SCALING EXPERIMENT")
	print("="*70)

	data_dir = f"{DATA_BASE}/icews_wiki"
	aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
	ref_pairs = os.path.join(data_dir, "ref_pairs")

	# Read all reference pairs for ground truth
	ref_set = load_ref_pairs(ref_pairs)
	kg1_entities = list(set(p[0] for p in ref_set))

	# Important entities file
	important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
	bak = important_file + ".bak_full"
	if not os.path.exists(bak):
	shutil.copy(important_file, bak)

	results = {}

	warmup_sizes = [100, 200, 500, 1000, 2000]

	for n_warmup in warmup_sizes:
	print(f"\n--- N_warmup = {n_warmup} ---")

	# Clear evorm state
	evorm_state = os.path.join(data_dir, "evorm_state")
	if os.path.exists(evorm_state):
	shutil.rmtree(evorm_state)

	# Clear aligned_entities
	if os.path.exists(aligned_file):
	os.remove(aligned_file)

	# Create subset
	with open(important_file, 'w') as f:
	for eid in kg1_entities[:n_warmup]:
	f.write(f"{eid}\n")

	import tokens_cal
	tokens_cal.global_tokens = 0

	from LLM1_label_selector import align_entities

	start_time = time.time()
	aligned_pairs = align_entities(data_dir, from_m3=False)
	elapsed = time.time() - start_time

	tokens = tokens_cal.global_tokens
	tokens_cal.global_tokens = 0

	metrics = compute_metrics(aligned_file, ref_pairs)

	results[f'n_warmup_{n_warmup}'] = {
	'aligned_pairs': len(aligned_pairs),
	'time_seconds': elapsed,
	'tokens': tokens,
	'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
	'metrics': metrics,
	}

	print(f" Aligned: {len(aligned_pairs)}, Tokens: {tokens}, "
	f"Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
	print(f" Hits@1: {metrics.get('Hits@1', 'N/A')}")

	# Restore
	if os.path.exists(bak):
	shutil.copy(bak, important_file)

	# Save results
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	result_file = os.path.join(RESULTS_DIR, f"coldstart_{timestamp}.json")
	ensure_dir(RESULTS_DIR)
	with open(result_file, 'w') as f:
	json.dump(results, f, indent=2)

	print(f"\nCold-start results saved to: {result_file}")

	# Print summary
	print("\n" + "="*70)
	print("COLD-START SCALING SUMMARY")
	print("="*70)
	print(f"{'N_warmup':<12} {'Hits@1':<10} {'Tokens':<12} {'Tokens/Pair':<12} {'Time(s)':<10}")
	print("-"*56)
	for n_warmup in warmup_sizes:
	key = f'n_warmup_{n_warmup}'
	if key in results:
	r = results[key]
	m = r['metrics']
	print(f"{n_warmup:<12} {m.get('Hits@1', 0):<10.4f} {r['tokens']:<12} {r['tokens_per_pair']:<12.1f} {r['time_seconds']:<10.1f}")

	return results

	def run_efficiency_analysis():
	"""RQ4: Efficiency Analysis - token consumption and latency."""
	print("\n" + "="*70)
	print("RQ4: EFFICIENCY ANALYSIS")
	print("="*70)

	data_dir = f"{DATA_BASE}/icews_wiki"
	aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
	ref_pairs = os.path.join(data_dir, "ref_pairs")

	important_file = os.path.join(data_dir, "message_pool", "important_entities.txt")
	bak = important_file + ".bak_full"
	if not os.path.exists(bak):
	shutil.copy(important_file, bak)

	results = {}

	# Test with different entity counts
	entity_counts = [100, 200, 500]

	for n_entities in entity_counts:
	print(f"\n--- N_entities = {n_entities} ---")

	# Clear evorm state
	evorm_state = os.path.join(data_dir, "evorm_state")
	if os.path.exists(evorm_state):
	shutil.rmtree(evorm_state)

	if os.path.exists(aligned_file):
	os.remove(aligned_file)

	# Subset
	ref_set = load_ref_pairs(ref_pairs)
	kg1_entities = list(set(p[0] for p in ref_set))
	with open(important_file, 'w') as f:
	for eid in kg1_entities[:n_entities]:
	f.write(f"{eid}\n")

	import tokens_cal
	tokens_cal.global_tokens = 0

	from LLM1_label_selector import align_entities

	# Measure latency
	latencies = []
	original_task = None

	start_time = time.time()
	aligned_pairs = align_entities(data_dir, from_m3=False)
	total_time = time.time() - start_time

	tokens = tokens_cal.global_tokens
	tokens_cal.global_tokens = 0

	metrics = compute_metrics(aligned_file, ref_pairs)

	n_pairs = len(aligned_pairs)
	results[f'n_{n_entities}'] = {
	'n_pairs': n_pairs,
	'total_time_s': total_time,
	'time_per_pair_s': total_time / max(1, n_pairs),
	'tokens': tokens,
	'tokens_per_pair': tokens / max(1, n_pairs),
	'metrics': metrics,
	}

	print(f" Pairs: {n_pairs}, Time: {total_time:.1f}s, "
	f"Time/Pair: {total_time/max(1,n_pairs):.2f}s, "
	f"Tokens/Pair: {tokens/max(1,n_pairs):.1f}")

	# Restore
	if os.path.exists(bak):
	shutil.copy(bak, important_file)

	# Save
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	result_file = os.path.join(RESULTS_DIR, f"efficiency_{timestamp}.json")
	ensure_dir(RESULTS_DIR)
	with open(result_file, 'w') as f:
	json.dump(results, f, indent=2)

	print(f"\nEfficiency results saved to: {result_file}")

	# Summary
	print("\n" + "="*70)
	print("EFFICIENCY SUMMARY")
	print("="*70)
	print(f"{'N_Entities':<12} {'Pairs':<8} {'Time(s)':<10} {'Time/Pair':<12} {'Tokens/Pair':<12} {'Hits@1':<10}")
	print("-"*64)
	for n_entities in entity_counts:
	key = f'n_{n_entities}'
	if key in results:
	r = results[key]
	m = r['metrics']
	print(f"{n_entities:<12} {r['n_pairs']:<8} {r['total_time_s']:<10.1f} {r['time_per_pair_s']:<12.2f} {r['tokens_per_pair']:<12.1f} {m.get('Hits@1', 0):<10.4f}")

	return results

	def run_main_experiment():
	"""RQ1: Main Results - full dataset run."""
	print("\n" + "="*70)
	print("RQ1: MAIN EXPERIMENT - ICEWS-WIKI (Full Dataset)")
	print("="*70)

	data_dir = f"{DATA_BASE}/icews_wiki"
	aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
	ref_pairs = os.path.join(data_dir, "ref_pairs")

	# Clear evorm state
	evorm_state = os.path.join(data_dir, "evorm_state")
	if os.path.exists(evorm_state):
	shutil.rmtree(evorm_state)

	if os.path.exists(aligned_file):
	os.remove(aligned_file)

	import tokens_cal
	tokens_cal.global_tokens = 0

	from LLM1_label_selector import align_entities

	start_time = time.time()
	aligned_pairs = align_entities(data_dir, from_m3=False)
	elapsed = time.time() - start_time

	tokens = tokens_cal.global_tokens
	tokens_cal.global_tokens = 0

	metrics = compute_metrics(aligned_file, ref_pairs)

	results = {
	'dataset': 'ICEWS-WIKI',
	'n_ref_pairs': 3540,
	'aligned_pairs': len(aligned_pairs),
	'time_seconds': elapsed,
	'tokens': tokens,
	'tokens_per_pair': tokens / max(1, len(aligned_pairs)),
	'metrics': metrics,
	}

	# Check evorm state
	if os.path.exists(evorm_state):
	rules_file = os.path.join(evorm_state, "rules.json")
	if os.path.exists(rules_file):
	with open(rules_file) as f:
	rules = json.load(f)
	results['n_rules'] = len(rules)
	hyperedges_file = os.path.join(evorm_state, "hyperedges.json")
	if os.path.exists(hyperedges_file):
	with open(hyperedges_file) as f:
	hedges = json.load(f)
	results['n_hyperedges'] = len(hedges)

	# Save
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	result_file = os.path.join(RESULTS_DIR, f"main_icews_wiki_{timestamp}.json")
	ensure_dir(RESULTS_DIR)
	with open(result_file, 'w') as f:
	json.dump(results, f, indent=2)

	print(f"Main results saved to: {result_file}")
	print(f"\n Dataset: ICEWS-WIKI")
	print(f" Aligned: {len(aligned_pairs)}, Time: {elapsed:.1f}s")
	print(f" Tokens: {tokens}, Tokens/Pair: {tokens/max(1,len(aligned_pairs)):.1f}")
	print(f" Hits@1: {metrics.get('Hits@1', 'N/A')}")
	print(f" MRR: {metrics.get('MRR', 'N/A')}")
	print(f" Precision: {metrics.get('Precision', 'N/A')}")
	print(f" Recall: {metrics.get('Recall', 'N/A')}")
	if 'n_rules' in results:
	print(f" Rules: {results['n_rules']}, Hyperedges: {results.get('n_hyperedges', 'N/A')}")

	return results

	def generate_latex_table(all_results):
	"""Generate paper-style LaTeX table."""
	timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
	table_file = os.path.join(RESULTS_DIR, f"paper_table_{timestamp}.tex")

	latex = r"""\begin{table}[t]
	\centering
	\caption{Experimental Results of EvoRM on ICEWS-WIKI}
	\label{tab:main_results}
	\begin{tabular}{lcccc}
	\hline
	\textbf{Method} & \textbf{Hits@1} & \textbf{MRR} & \textbf{Precision} & \textbf{Recall} \\
	\hline
	"""

	# Main results
	if 'main' in all_results:
	m = all_results['main']['metrics']
	latex += f"EvoRM (Ours) & {m['Hits@1']:.4f} & {m['MRR']:.4f} & {m['Precision']:.4f} & {m['Recall']:.4f} \\\\\n"

	latex += r"""\hline
	\end{tabular}
	\end{table}
	"""

	with open(table_file, 'w') as f:
	f.write(latex)

	print(f"LaTeX table saved to: {table_file}")
	return table_file

	def main():
	parser = argparse.ArgumentParser(description='EvoRM Comprehensive Experiment Runner')
	parser.add_argument('--all', action='store_true', help='Run all experiments')
	parser.add_argument('--ablation', action='store_true', help='Run ablation study (RQ3)')
	parser.add_argument('--coldstart', action='store_true', help='Run cold-start scaling (RQ5)')
	parser.add_argument('--main', action='store_true', help='Run main experiment (RQ1)')
	parser.add_argument('--efficiency', action='store_true', help='Run efficiency analysis (RQ4)')
	args = parser.parse_args()

	# Default: run everything if no flag specified
	run_all = args.all or not any([args.ablation, args.coldstart, args.main, args.efficiency])

	all_results = {}

	if run_all or args.main:
	all_results['main'] = run_main_experiment()

	if run_all or args.ablation:
	all_results['ablation'] = run_ablation_study()

	if run_all or args.coldstart:
	all_results['coldstart'] = run_coldstart_experiment()

	if run_all or args.efficiency:
	all_results['efficiency'] = run_efficiency_analysis()

	# Generate paper table
	generate_latex_table(all_results)

	print("\n" + "="*70)
	print("ALL EXPERIMENTS COMPLETE!")
	print("="*70)
	print(f"Results saved in: {RESULTS_DIR}")
	for f in sorted(os.listdir(RESULTS_DIR)):
	print(f" {f}")

	if __name__ == '__main__':
	main()