#!/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()