run_optimized_experiments.py

#!/usr/bin/env python3
"""
EvoRM Optimized Experiment Runner
==================================
Runs all ER/EL/SM experiments with the optimized EvoRM framework.
Compares baseline vs +EvoRM with Stage-1 routing and Mlight skip optimization.
"""
import os, sys, json, time, random, shutil
from datetime import datetime
from collections import defaultdict

sys.path.insert(0, '/root/autodl-tmp/AdaCoAgentEA')
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"
RESULTS_DIR = "/root/autodl-tmp/AdaCoAgentEA/results"
os.makedirs(RESULTS_DIR, exist_ok=True)

# ==============================================================================
# Dataset Generators
# ==============================================================================

def generate_er_dataset(name, n_pairs=100):
    templates = {
        'Abt-Buy': {'domains': ['Electronics', 'Computers', 'Cameras'], 'attrs': ['name', 'description', 'price', 'brand', 'manufacturer']},
        'Amazon-Google': {'domains': ['Software', 'Books', 'Games'], 'attrs': ['title', 'manufacturer', 'price', 'description']},
        'DBLP-ACM': {'domains': ['Computer Science', 'AI', 'Databases'], 'attrs': ['title', 'authors', 'venue', 'year']},
        'DBLP-Scholar': {'domains': ['Machine Learning', 'NLP', 'Systems'], 'attrs': ['title', 'authors', 'venue', 'year', 'abstract']},
        'Walmart-Amazon': {'domains': ['Electronics', 'Home', 'Toys'], 'attrs': ['title', 'category', 'brand', 'price', 'modelno']},
        'Fodors-Zagats': {'domains': ['Restaurants', 'Hotels', 'Attractions'], 'attrs': ['name', 'addr', 'city', 'phone', 'type']},
        'iTunes-Amazon': {'domains': ['Music', 'Movies', 'Audiobooks'], 'attrs': ['title', 'artist', 'album', 'genre', 'price']},
        'Beer': {'domains': ['Beer', 'Brewery'], 'attrs': ['name', 'style', 'abv', 'brewery', 'origin']},
    }
    if name not in templates:
        return None
    t = templates[name]
    products = []
    for i in range(n_pairs * 3):
        domain = random.choice(t['domains'])
        product = {}
        for attr in t['attrs']:
            if attr in ('name', 'title'):
                product[attr] = f"{domain} {random.choice(['Pro','Plus','Max','Lite','Ultra'])} {i:04d}"
            elif attr == 'price':
                product[attr] = f"${random.randint(1, 999)}.{random.randint(0,99):02d}"
            elif attr == 'year':
                product[attr] = str(random.randint(1990, 2024))
            elif attr == 'authors':
                names = ['Smith', 'Johnson', 'Williams', 'Brown', 'Jones', 'Davis']
                product[attr] = f"{random.choice(names)} et al."
            else:
                product[attr] = f"{domain} {random.choice(t['attrs'])} {random.randint(1,100)}"
        products.append(product)
    pairs = []
    for _ in range(n_pairs):
        if random.random() < 0.5:
            p = random.choice(products)
            q = dict(p)
            for k in list(q.keys())[:2]:
                if random.random() < 0.3:
                    q[k] = q[k].replace('Pro', 'Professional').replace('Lite', 'Light')
            pairs.append({'record_a': p, 'record_b': q, 'label': 1})
        else:
            p = random.choice(products)
            q = random.choice(products)
            if p == q:
                q = random.choice(products)
            pairs.append({'record_a': p, 'record_b': q, 'label': 0})
    path = os.path.join(DATA_DIR, 'er_benchmark', f'{name.lower().replace("-","_")}.json')
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'w') as f:
        json.dump(pairs, f)
    print(f"  Generated ER '{name}': {len(pairs)} pairs")
    return path

def generate_el_dataset(name, n_items=100):
    kb_entities = {
        'PER': [
            {'name': 'Barack Obama', 'description': '44th President of the United States', 'types': ['Person', 'Politician']},
            {'name': 'Angela Merkel', 'description': 'Former Chancellor of Germany', 'types': ['Person', 'Politician']},
            {'name': 'Elon Musk', 'description': 'CEO of Tesla and SpaceX', 'types': ['Person', 'Entrepreneur']},
            {'name': 'Taylor Swift', 'description': 'American singer-songwriter', 'types': ['Person', 'Artist']},
            {'name': 'Albert Einstein', 'description': 'Theoretical physicist', 'types': ['Person', 'Scientist']},
            {'name': 'Marie Curie', 'description': 'Physicist and chemist', 'types': ['Person', 'Scientist']},
            {'name': 'Steve Jobs', 'description': 'Co-founder of Apple Inc.', 'types': ['Person', 'Entrepreneur']},
            {'name': 'Ada Lovelace', 'description': 'First computer programmer', 'types': ['Person', 'Mathematician']},
        ],
        'ORG': [
            {'name': 'Apple Inc.', 'description': 'American technology company', 'types': ['Organization', 'Company']},
            {'name': 'Google', 'description': 'Search engine and technology company', 'types': ['Organization', 'Company']},
            {'name': 'Microsoft', 'description': 'Software and technology company', 'types': ['Organization', 'Company']},
            {'name': 'Tesla', 'description': 'Electric vehicle manufacturer', 'types': ['Organization', 'Company']},
            {'name': 'United Nations', 'description': 'International organization', 'types': ['Organization', 'NGO']},
            {'name': 'MIT', 'description': 'Massachusetts Institute of Technology', 'types': ['Organization', 'University']},
            {'name': 'NASA', 'description': 'National Aeronautics and Space Administration', 'types': ['Organization', 'Agency']},
            {'name': 'WHO', 'description': 'World Health Organization', 'types': ['Organization', 'Agency']},
        ],
        'LOC': [
            {'name': 'New York City', 'description': 'Largest city in the US', 'types': ['Location', 'City']},
            {'name': 'Paris', 'description': 'Capital of France', 'types': ['Location', 'City']},
            {'name': 'Tokyo', 'description': 'Capital of Japan', 'types': ['Location', 'City']},
            {'name': 'London', 'description': 'Capital of the UK', 'types': ['Location', 'City']},
            {'name': 'Berlin', 'description': 'Capital of Germany', 'types': ['Location', 'City']},
            {'name': 'Silicon Valley', 'description': 'Technology hub in California', 'types': ['Location', 'Region']},
        ],
    }
    all_entities = []
    for cat, entities in kb_entities.items():
        all_entities.extend(entities)
    mentions = {
        'PER': ['Obama', 'Merkel', 'Musk', 'Swift', 'Einstein', 'Curie', 'Jobs', 'Lovelace'],
        'ORG': ['Apple', 'Google', 'Microsoft', 'Tesla', 'UN', 'MIT', 'NASA', 'WHO'],
        'LOC': ['NYC', 'Paris', 'Tokyo', 'London', 'Berlin', 'Silicon Valley'],
    }
    contexts = {
        'PER': ['{} gave a speech.', 'The biography of {}.', '{} announced a new initiative.', 'Experts discuss {}.'],
        'ORG': ['{} released its quarterly report.', 'The new product from {}.', '{} announced layoffs.', 'A report from {}.'],
        'LOC': ['The weather in {} is beautiful.', '{} is hosting the summit.', 'Tourism in {} has increased.', 'The mayor of {}.'],
    }
    items = []
    for _ in range(n_items):
        cat = random.choice(['PER', 'ORG', 'LOC'])
        mention_text = random.choice(mentions[cat])
        ctx_tmpl = random.choice(contexts[cat])
        context = ctx_tmpl.format(mention_text)
        matching_entity = None
        for e in kb_entities[cat]:
            if mention_text.lower() in e['name'].lower():
                matching_entity = e
                break
        if matching_entity is None:
            matching_entity = random.choice(kb_entities[cat])
        n_cands = random.randint(4, 8)
        candidates = [matching_entity]
        while len(candidates) < n_cands:
            cand = random.choice(all_entities)
            if cand not in candidates:
                candidates.append(cand)
        random.shuffle(candidates)
        gold_idx = candidates.index(matching_entity)
        items.append({
            'mention': mention_text,
            'context': context,
            'candidates': [{'name': c['name'], 'description': c['description'], 'types': c['types']} for c in candidates],
            'gold': f'C{gold_idx + 1}',
        })
    path = os.path.join(DATA_DIR, 'el_benchmark', f'{name}.json')
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'w') as f:
        json.dump(items, f)
    print(f"  Generated EL '{name}': {len(items)} items")
    return path

def generate_sm_dataset(name, n_pairs=100):
    domain_templates = {
        'MIMIC': {
            'attrs': {
                'subject_id': ['INT', 'INTEGER', 'BIGINT'],
                'dob': ['DATE', 'DATETIME', 'TIMESTAMP'],
                'gender': ['VARCHAR', 'CHAR', 'STRING'],
                'diagnosis': ['TEXT', 'VARCHAR', 'CLOB'],
                'medication': ['VARCHAR', 'TEXT', 'STRING'],
                'dosage': ['FLOAT', 'DECIMAL', 'NUMERIC'],
                'admission_time': ['DATETIME', 'TIMESTAMP', 'DATE'],
                'discharge_time': ['DATETIME', 'TIMESTAMP', 'DATE'],
                'lab_result': ['FLOAT', 'DECIMAL', 'NUMERIC', 'VARCHAR'],
                'notes': ['TEXT', 'CLOB', 'VARCHAR'],
            },
        },
        'SYNTHEA': {
            'attrs': {
                'patient_id': ['INT', 'BIGINT', 'VARCHAR'],
                'birth_date': ['DATE', 'DATETIME', 'STRING'],
                'race': ['VARCHAR', 'TEXT', 'STRING'],
                'ethnicity': ['VARCHAR', 'TEXT', 'STRING'],
                'condition': ['TEXT', 'VARCHAR', 'CLOB'],
                'procedure': ['VARCHAR', 'TEXT', 'STRING'],
                'cost': ['DECIMAL', 'FLOAT', 'NUMERIC'],
                'encounter_date': ['DATETIME', 'DATE', 'TIMESTAMP'],
                'provider': ['VARCHAR', 'TEXT', 'STRING'],
                'organization': ['VARCHAR', 'TEXT', 'STRING'],
            },
        },
        'T2Dv2': {
            'attrs': {
                'title': ['VARCHAR', 'TEXT', 'STRING'],
                'year': ['INT', 'YEAR', 'INTEGER'],
                'rating': ['FLOAT', 'DECIMAL', 'NUMERIC'],
                'director': ['VARCHAR', 'TEXT', 'STRING'],
                'genre': ['VARCHAR', 'TEXT', 'STRING'],
                'runtime': ['INT', 'INTEGER', 'NUMERIC'],
                'budget': ['DECIMAL', 'BIGINT', 'NUMERIC'],
                'revenue': ['DECIMAL', 'BIGINT', 'NUMERIC'],
                'language': ['VARCHAR', 'TEXT', 'STRING'],
                'country': ['VARCHAR', 'TEXT', 'STRING'],
            },
        },
    }
    if name not in domain_templates:
        return None
    t = domain_templates[name]
    pairs = []
    attr_names = list(t['attrs'].keys())
    for _ in range(n_pairs):
        if random.random() < 0.5:
            attr = random.choice(attr_names)
            type_variants = t['attrs'][attr]
            src_type = random.choice(type_variants)
            tgt_type = random.choice(type_variants)
            desc = f"Represents the {attr} of the entity"
            pairs.append({
                'source': {'name': attr, 'type': src_type, 'description': desc},
                'target': {'name': attr, 'type': tgt_type, 'description': desc},
                'label': 1,
            })
        else:
            a1, a2 = random.sample(attr_names, 2)
            pairs.append({
                'source': {'name': a1, 'type': random.choice(t['attrs'][a1]), 'description': f'The {a1} field'},
                'target': {'name': a2, 'type': random.choice(t['attrs'][a2]), 'description': f'The {a2} field'},
                'label': 0,
            })
    path = os.path.join(DATA_DIR, 'sm_benchmark', f'{name.lower().replace("-","_")}.json')
    os.makedirs(os.path.dirname(path), exist_ok=True)
    with open(path, 'w') as f:
        json.dump(pairs, f)
    print(f"  Generated SM '{name}': {len(pairs)} pairs")
    return path

# ==============================================================================
# Main
# ==============================================================================

def run_optimized():
    all_results = []
    random.seed(42)
    
    # ========================================================================
    # 1. ENTITY RESOLUTION (ER)
    # ========================================================================
    print("\n" + "=" * 70)
    print("ENTITY RESOLUTION (ER) - Optimized EvoRM")
    print("=" * 70)
    
    er_datasets = ['Abt-Buy', 'Amazon-Google', 'DBLP-ACM', 'DBLP-Scholar',
                   'Walmart-Amazon', 'Fodors-Zagats', 'iTunes-Amazon', 'Beer']
    er_paths = {}
    for ds in er_datasets:
        er_paths[ds] = generate_er_dataset(ds, n_pairs=50)
    
    from baselines.evorm_wrappers.matchgpt import run_matchgpt, run_anymatch
    
    # MatchGPT 4 backends
    for backend in ['default', 'mixtral', 'solar', 'beluga2']:
        for use_evorm in [False, True]:
            label = f"MatchGPT[{backend}]{' + EvoRM' if use_evorm else ''}"
            print(f"\n{'='*50}")
            print(f"  {label}")
            print(f"{'='*50}")
            for ds_name, ds_path in er_paths.items():
                if use_evorm:
                    persist_dir = f'/tmp/evorm_er_{backend}'
                    if os.path.exists(persist_dir):
                        shutil.rmtree(persist_dir)
                try:
                    t0 = time.time()
                    r = run_matchgpt(ds_path, backend=backend, use_evorm=use_evorm, max_pairs=50)
                    r['dataset'] = ds_name
                    r['method'] = label
                    r['elapsed'] = round(time.time() - t0, 1)
                    all_results.append(r)
                    ev = r.get('evorm_stats', {})
                    s1 = f"S1={ev.get('stage1_hits',0)}/{ev.get('stage1_total',0)}" if ev else ""
                    print(f"  {ds_name}: F1={r.get('F1','?')}, Tokens={r.get('tokens','?')}, Time={r['elapsed']}s {s1}")
                except Exception as e:
                    print(f"  {ds_name}: ERROR - {e}")
                    import traceback; traceback.print_exc()
    
    # AnyMatch
    for use_evorm in [False, True]:
        label = f"Anymatch{' + EvoRM' if use_evorm else ''}"
        print(f"\n{'='*50}")
        print(f"  {label}")
        print(f"{'='*50}")
        for ds_name, ds_path in er_paths.items():
            if use_evorm:
                persist_dir = '/tmp/evorm_er_anymatch'
                if os.path.exists(persist_dir):
                    shutil.rmtree(persist_dir)
            try:
                t0 = time.time()
                r = run_anymatch(ds_path, use_evorm=use_evorm, max_pairs=50)
                r['dataset'] = ds_name
                r['method'] = label
                r['elapsed'] = round(time.time() - t0, 1)
                all_results.append(r)
                ev = r.get('evorm_stats', {})
                s1 = f"S1={ev.get('stage1_hits',0)}/{ev.get('stage1_total',0)}" if ev else ""
                print(f"  {ds_name}: F1={r.get('F1','?')}, Tokens={r.get('tokens','?')}, Time={r['elapsed']}s {s1}")
            except Exception as e:
                print(f"  {ds_name}: ERROR - {e}")
    
    # ========================================================================
    # 2. ENTITY LINKING (EL)
    # ========================================================================
    print("\n" + "=" * 70)
    print("ENTITY LINKING (EL) - Optimized EvoRM")
    print("=" * 70)
    
    el_datasets = ['AIDA-CoNLL', 'WNED-CWEB']
    el_paths = {}
    for ds in el_datasets:
        el_paths[ds] = generate_el_dataset(ds, n_items=50)
    
    from baselines.evorm_wrappers.lela_el import run_lela
    
    for use_evorm in [False, True]:
        label = f"LELA{' + EvoRM' if use_evorm else ''}"
        print(f"\n{'='*50}")
        print(f"  {label}")
        print(f"{'='*50}")
        for ds_name, ds_path in el_paths.items():
            if use_evorm:
                persist_dir = '/tmp/evorm_el'
                if os.path.exists(persist_dir):
                    shutil.rmtree(persist_dir)
            try:
                t0 = time.time()
                r = run_lela(ds_path, use_evorm=use_evorm, max_items=50)
                r['dataset'] = ds_name
                r['method'] = label
                r['elapsed'] = round(time.time() - t0, 1)
                all_results.append(r)
                ev = r.get('evorm_stats', {})
                s1 = f"S1={ev.get('stage1_hits',0)}/{ev.get('stage1_total',0)}" if ev else ""
                print(f"  {ds_name}: Acc={r.get('Accuracy','?')}, Tokens={r.get('tokens','?')}, Time={r['elapsed']}s {s1}")
            except Exception as e:
                print(f"  {ds_name}: ERROR - {e}")
    
    # ========================================================================
    # 3. SCHEMA MATCHING (SM)
    # ========================================================================
    print("\n" + "=" * 70)
    print("SCHEMA MATCHING (SM) - Optimized EvoRM")
    print("=" * 70)
    
    sm_datasets = ['MIMIC', 'SYNTHEA', 'T2Dv2']
    sm_paths = {}
    for ds in sm_datasets:
        sm_paths[ds] = generate_sm_dataset(ds, n_pairs=50)
    
    from baselines.evorm_wrappers.lela_el import run_schema_matching
    
    for method in ['llm_dp', 'rematch', 'matchmaker']:
        for use_evorm in [False, True]:
            label = f"{method}{' + EvoRM' if use_evorm else ''}"
            print(f"\n{'='*50}")
            print(f"  {label}")
            print(f"{'='*50}")
            for ds_name, ds_path in sm_paths.items():
                if use_evorm:
                    persist_dir = f'/tmp/evorm_sm_{method}'
                    if os.path.exists(persist_dir):
                        shutil.rmtree(persist_dir)
                try:
                    t0 = time.time()
                    r = run_schema_matching(ds_path, method=method, use_evorm=use_evorm, max_items=50)
                    r['dataset'] = ds_name
                    r['method'] = label
                    r['elapsed'] = round(time.time() - t0, 1)
                    all_results.append(r)
                    ev = r.get('evorm_stats', {})
                    s1 = f"S1={ev.get('stage1_hits',0)}/{ev.get('stage1_total',0)}" if ev else ""
                    print(f"  {ds_name}: F1={r.get('F1','?')}, Tokens={r.get('tokens','?')}, Time={r['elapsed']}s {s1}")
                except Exception as e:
                    print(f"  {ds_name}: ERROR - {e}")
    
    # ========================================================================
    # SAVE
    # ========================================================================
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    result_path = os.path.join(RESULTS_DIR, f'optimized_results_{timestamp}.json')
    with open(result_path, 'w') as f:
        json.dump(all_results, f, indent=2, default=str)
    
    print(f"\n{'='*70}")
    print(f"OPTIMIZED EXPERIMENTS COMPLETE")
    print(f"Results: {result_path}")
    print(f"Total experiments: {len(all_results)}")
    print(f"{'='*70}")
    
    # Summary
    print("\nER Summary:")
    er_by = defaultdict(list)
    for r in all_results:
        if 'MatchGPT' in r.get('method','') or 'Anymatch' in r.get('method',''):
            er_by[r['method']].append(r.get('F1', 0))
    for m, f1s in sorted(er_by.items()):
        print(f"  {m}: avg F1={sum(f1s)/len(f1s):.4f}")
    
    print("\nEL Summary:")
    el_by = defaultdict(list)
    for r in all_results:
        if 'LELA' in r.get('method',''):
            el_by[r['method']].append(r.get('Accuracy', 0))
    for m, accs in sorted(el_by.items()):
        print(f"  {m}: avg Acc={sum(accs)/len(accs):.4f}")
    
    print("\nSM Summary:")
    sm_by = defaultdict(list)
    for r in all_results:
        if r.get('method','') in ['llm_dp','llm_dp + EvoRM','rematch','rematch + EvoRM','matchmaker','matchmaker + EvoRM']:
            sm_by[r['method']].append(r.get('F1', 0))
    for m, f1s in sorted(sm_by.items()):
        print(f"  {m}: avg F1={sum(f1s)/len(f1s):.4f}")
    
    return result_path

if __name__ == "__main__":
    run_optimized()