code/experiments/run_baselines_v2.py

#!/usr/bin/env python3
"""
EA Baselines Comparison - All baselines with and without EvoRM
Uses ref_pairs order for entity selection.
"""
import os, sys, json, time, shutil
from datetime import datetime

sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')
sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/baselines/evorm_wrappers')

os.environ["OPENAI_API_BASE"] = "https://hk.xty.app/v1"
os.environ["OPENAI_API_KEY"] = "sk-7a7Ev4VcVyysPLT5hqtqIVD6PybzJ1ZlEIVZddIR3NtZvPgK"

DATA_DIR = "/root/autodl-tmp/AdaCoAgentEA/data/icews_wiki"
RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
ALIGNED_FILE = os.path.join(DATA_DIR, "message_pool", "aligned_entities.txt")
REF_PAIRS_FILE = os.path.join(DATA_DIR, "ref_pairs")
IMPORTANT_FILE = os.path.join(DATA_DIR, "message_pool", "important_entities.txt")
RETRIEVER_FILE = os.path.join(DATA_DIR, "message_pool", "retriever_outputs.txt")
EVORM_STATE = os.path.join(DATA_DIR, "evorm_state")

os.makedirs(RESULTS_DIR, exist_ok=True)

def load_ref_pairs():
    pairs = {}
    with open(REF_PAIRS_FILE, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    pairs[int(parts[0])] = int(parts[1])
                except:
                    pass
    return pairs

def compute_metrics(ref_subset=None):
    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:
                        continue
    if ref_subset:
        ref_set = set(ref_subset.items())
    else:
        ref_set = set(load_ref_pairs().items())
    if not ref_set:
        return {'error': 'No ref pairs'}
    tp = len(aligned & ref_set)
    fp = len(aligned - ref_set)
    fn = len(ref_set - aligned)
    return {
        'Hits@1': round(tp / max(1, len(ref_set)), 4),
        'MRR': round(tp / max(1, len(ref_set)), 4),
        'Precision': round(tp / max(1, tp + fp), 4),
        'Recall': round(tp / max(1, tp + fn), 4),
        'F1': round(2 * tp / max(1, 2*tp + fp + fn), 4),
        'TP': tp, 'FP': fp, 'FN': fn,
        'Total_Ref': len(ref_set),
        'Total_Aligned': len(aligned),
    }

def clear_state():
    if os.path.exists(EVORM_STATE):
        shutil.rmtree(EVORM_STATE)
    if os.path.exists(ALIGNED_FILE):
        os.remove(ALIGNED_FILE)

def setup_subset(n_entities, ref_map):
    """Create subset using ref_pairs order."""
    ref_order = []
    with open(REF_PAIRS_FILE, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    e1 = int(parts[0])
                    if e1 not in ref_order:
                        ref_order.append(e1)
                except:
                    pass
    actual_n = min(n_entities, len(ref_order))
    subset_kg1 = set(ref_order[:actual_n])
    subset_lines = []
    with open(RETRIEVER_FILE, 'r') as f:
        for line in f:
            parts = line.strip().split('\t')
            if len(parts) == 2:
                try:
                    if int(parts[0]) in subset_kg1:
                        subset_lines.append(line)
                except:
                    pass
    with open(IMPORTANT_FILE, 'w') as f:
        f.writelines(subset_lines)
    subset_ref = {e1: ref_map[e1] for e1 in subset_kg1 if e1 in ref_map}
    return subset_ref

def run_adacoagent_baseline(data_dir, use_evorm=False):
    """Run AdaCoAgentEA (baseline or +EvoRM)."""
    clear_state()
    import tokens_cal
    tokens_cal.global_tokens = 0
    t0 = time.time()
    if use_evorm:
        from LLM1_label_selector import align_entities
    else:
        from LLM1_label_selector_baseline import align_entities
    aligned = align_entities(data_dir, from_m3=False)
    elapsed = time.time() - t0
    return aligned, elapsed, tokens_cal.global_tokens

def main():
    n_entities = 200  # Subset size for comparison
    ref_map = load_ref_pairs()
    print(f"Loaded {len(ref_map)} reference pairs")
    
    subset_ref = setup_subset(n_entities, ref_map)
    print(f"Subset: {len(subset_ref)} entities with ref pairs")
    
    # Baselines to run
    baselines = [
        ('AdaCoAgentEA', 'adacoagent'),
        ('ZeroCoT', 'zerocot'),
        ('Self-Consistency', 'self_consistency'),
        ('ChatEA', 'chatea'),
        ('Collaboration-Hard', 'cohard'),
    ]
    
    import tokens_cal
    
    results = []
    for bl_name, bl_key in baselines:
        for use_evorm in [False, True]:
            mode = f"{bl_name}{' + EvoRM' if use_evorm else ''}"
            print(f"\n{'='*50}")
            print(f"  {mode}")
            print(f"{'='*50}")
            
            clear_state()
            setup_subset(n_entities, ref_map)  # Re-setup after clear_state
            
            try:
                tokens_cal.global_tokens = 0
                t0 = time.time()
                
                if bl_key == 'adacoagent':
                    if use_evorm:
                        from LLM1_label_selector import align_entities
                    else:
                        from LLM1_label_selector_baseline import align_entities
                    aligned = align_entities(DATA_DIR, from_m3=False)
                elif bl_key == 'zerocot':
                    from zero_cot import align_entities as zc
                    aligned = zc(DATA_DIR, use_evorm=use_evorm)
                elif bl_key == 'self_consistency':
                    from self_consistency import align_entities as sc
                    aligned = sc(DATA_DIR, use_evorm=use_evorm)
                elif bl_key == 'chatea':
                    from chat_ea import align_entities as ce
                    aligned = ce(DATA_DIR, use_evorm=use_evorm)
                elif bl_key == 'cohard':
                    from cohard import align_entities as ch
                    aligned = ch(DATA_DIR, use_evorm=use_evorm)
                else:
                    raise ValueError(f"Unknown: {bl_key}")
                
                elapsed = time.time() - t0
                tokens = tokens_cal.global_tokens
                metrics = compute_metrics(ref_subset=subset_ref)
                
                result = {
                    'method': mode,
                    'aligned_pairs': len(aligned),
                    'time_seconds': round(elapsed, 1),
                    'tokens': tokens,
                    'tokens_per_pair': round(tokens / max(1, len(aligned)), 1),
                    'metrics': metrics,
                }
                results.append(result)
                
                print(f"  Aligned: {len(aligned)}, Time: {elapsed:.1f}s, Tokens: {tokens}")
                print(f"  Hits@1: {metrics['Hits@1']}, Precision: {metrics['Precision']}, Recall: {metrics['Recall']}")
                
            except Exception as e:
                print(f"  FAILED: {e}")
                import traceback
                traceback.print_exc()
    
    # Save
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = os.path.join(RESULTS_DIR, f"baselines_comparison_{timestamp}.json")
    with open(result_file, 'w') as f:
        json.dump({
            'experiment': 'EA Baselines Comparison',
            'dataset': 'ICEWS-WIKI',
            'n_entities': n_entities,
            'results': results,
        }, f, indent=2)
    
    # Summary
    print(f"\n{'='*80}")
    print("BASELINES COMPARISON SUMMARY")
    print(f"{'='*80}")
    print(f"{'Method':<35} {'Hits@1':>8} {'Prec':>8} {'Recall':>8} {'F1':>8} {'Tokens':>10} {'Time':>8}")
    print("-" * 80)
    for r in results:
        m = r['metrics']
        print(f"{r['method']:<35} {m['Hits@1']:>8.4f} {m['Precision']:>8.4f} "
              f"{m['Recall']:>8.4f} {m['F1']:>8.4f} {r['tokens']:>10} {r['time_seconds']:>7.1f}s")
    
    print(f"\nResults: {result_file}")

if __name__ == '__main__':
    main()