Upload experiments/run_baselines_comparison.py with huggingface_hub

experiments/run_baselines_comparison.py

@@ -0,0 +1,187 @@
+#!/usr/bin/env python3
+"""
+EA Baselines + EvoRM Comparison Experiment
+===========================================
+Runs all EA baselines standalone and with EvoRM, computes comparison metrics.
+
+Baselines:
+  1. AdaCoAgentEA (baseline) - LLM1_label_selector_baseline.py
+  2. AdaCoAgentEA + EvoRM - LLM1_label_selector.py
+  3. ZeroCoT - Zero-shot Chain-of-Thought
+  4. ZeroCoT + EvoRM
+  5. Self-Consistency - Multi-sample voting
+  6. Self-Consistency + EvoRM
+  7. ChatEA - Chat-based multi-turn
+  8. ChatEA + EvoRM
+  9. Collaboration-Hard - Multi-agent voting
+  10. Collaboration-Hard + EvoRM
+
+Metrics: Hits@1, MRR, Precision, Recall, F1, Tokens, Latency
+"""
+
+import os, sys, json, time, shutil, argparse
+from datetime import datetime
+
+sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
+sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/baselines/evorm_wrappers')
+
+DATA_BASE = "/root/autodl-tmp/AdaCoAgentEA/data"
+RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
+os.makedirs(RESULTS_DIR, exist_ok=True)
+
+
+def load_ref_pairs(path):
+    pairs = set()
+    if not os.path.exists(path):
+        return pairs
+    with open(path, 'r') as f:
+        for line in f:
+            parts = line.strip().split('\t')
+            if len(parts) == 2:
+                try:
+                    pairs.add((int(parts[0]), int(parts[1])))
+                except ValueError:
+                    pass
+    return pairs
+
+
+def compute_metrics(aligned, ref_set):
+    aligned_set = set(aligned)
+    tp = len(aligned_set & ref_set)
+    fp = len(aligned_set - ref_set)
+    fn = len(ref_set - aligned_set)
+    precision = tp / max(1, tp + fp)
+    recall = tp / max(1, tp + fn)
+    f1 = 2 * precision * recall / max(0.001, precision + recall)
+    return {
+        'Hits@1': round(recall, 4),
+        'MRR': round(recall, 4),
+        'Precision': round(precision, 4),
+        'Recall': round(recall, 4),
+        'F1': round(f1, 4),
+        'TP': tp, 'FP': fp, 'FN': fn,
+        'Total_Ref': len(ref_set),
+        'Total_Aligned': len(aligned),
+    }
+
+
+def clear_state(data_dir):
+    aligned_file = os.path.join(data_dir, "message_pool", "aligned_entities.txt")
+    if os.path.exists(aligned_file):
+        os.remove(aligned_file)
+
+
+def run_baseline(baseline_name, data_dir, use_evorm=False, n_entities=None):
+    """Run a single baseline and return results."""
+    clear_state(data_dir)
+    import tokens_cal
+    tokens_cal.global_tokens = 0
+
+    t0 = time.time()
+
+    if baseline_name == 'adacoagent':
+        if use_evorm:
+            from Area2.LLM1_label_selector import align_entities
+            aligned = align_entities(data_dir, from_m3=False)
+        else:
+            sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA/Area2')
+            from LLM1_label_selector_baseline import align_entities
+            aligned = align_entities(data_dir, from_m3=False)
+    elif baseline_name == 'zerocot':
+        from zero_cot import align_entities as zc
+        aligned = zc(data_dir, use_evorm=use_evorm)
+    elif baseline_name == 'self_consistency':
+        from self_consistency import align_entities as sc
+        aligned = sc(data_dir, use_evorm=use_evorm)
+    elif baseline_name == 'chatea':
+        from chat_ea import align_entities as ce
+        aligned = ce(data_dir, use_evorm=use_evorm)
+    elif baseline_name == 'cohard':
+        from cohard import align_entities as ch
+        aligned = ch(data_dir, use_evorm=use_evorm)
+    else:
+        raise ValueError(f"Unknown baseline: {baseline_name}")
+
+    elapsed = time.time() - t0
+    tokens = tokens_cal.global_tokens
+    ref_pairs = load_ref_pairs(os.path.join(data_dir, "ref_pairs"))
+    metrics = compute_metrics(aligned, ref_pairs)
+
+    return {
+        'method': f"{baseline_name}{' + EvoRM' if use_evorm else ''}",
+        'aligned_pairs': len(aligned),
+        'time_seconds': round(elapsed, 1),
+        'tokens': tokens,
+        'tokens_per_pair': round(tokens / max(1, len(aligned)), 1),
+        'metrics': metrics,
+    }
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', default='icews_wiki')
+    parser.add_argument('--baselines', nargs='+', default=['adacoagent', 'zerocot', 'self_consistency', 'chatea', 'cohard'])
+    parser.add_argument('--n_entities', type=int, default=None)
+    args = parser.parse_args()
+
+    data_dir = os.path.join(DATA_BASE, args.dataset)
+    if not os.path.exists(data_dir):
+        print(f"ERROR: {data_dir} not found")
+        return
+
+    print(f"{'='*70}")
+    print(f"EA BASELINES COMPARISON: {args.dataset}")
+    print(f"{'='*70}")
+
+    results = []
+
+    for bl_name in args.baselines:
+        # Baseline (no EvoRM)
+        print(f"\n--- {bl_name} (baseline) ---")
+        try:
+            r = run_baseline(bl_name, data_dir, use_evorm=False)
+            results.append(r)
+            print(f"  Aligned: {r['aligned_pairs']}, Hits@1: {r['metrics']['Hits@1']}, "
+                  f"Tokens: {r['tokens']}, Time: {r['time_seconds']}s")
+        except Exception as e:
+            print(f"  FAILED: {e}")
+            import traceback; traceback.print_exc()
+
+        # +EvoRM
+        print(f"--- {bl_name} + EvoRM ---")
+        try:
+            r = run_baseline(bl_name, data_dir, use_evorm=True)
+            results.append(r)
+            print(f"  Aligned: {r['aligned_pairs']}, Hits@1: {r['metrics']['Hits@1']}, "
+                  f"Tokens: {r['tokens']}, Time: {r['time_seconds']}s")
+        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 + EvoRM Comparison',
+            'dataset': args.dataset,
+            'date': datetime.now().isoformat(),
+            'results': results,
+        }, f, indent=2, ensure_ascii=False)
+
+    # Summary table
+    print(f"\n{'='*70}")
+    print("SUMMARY")
+    print(f"{'='*70}")
+    print(f"{'Method':<35} {'Hits@1':>8} {'Precision':>10} {'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']:>10.4f} "
+              f"{m['F1']:>8.4f} {r['tokens']:>10} {r['time_seconds']:>7.1f}s")
+
+    print(f"\nResults: {result_file}")
+
+
+if __name__ == '__main__':
+    main()