benchmark_test.py

"""
Benchmark Test Script for Plain-English Translator
Compares SMC outputs from different models against Claude Opus 4.5 benchmarks.
"""

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import random
import json
from datetime import datetime

# ============================================================================
# CLAUDE OPUS 4.5 BENCHMARK TRANSLATIONS
# ============================================================================

BENCHMARKS = {
    "Legal": {
        "Force Majeure clause and why it might void our contract":
            "This is a 'nobody's fault' escape hatch. If something huge and uncontrollable happens—like a massive earthquake, a war, or a pandemic—neither of us can be blamed for not keeping our promises. It's like if you promised to meet a friend but a tornado blocked every road; you couldn't get there, but it wasn't your fault.",

        "Why we need to add an indemnification clause to protect your business":
            "This is a 'you cover me, I cover you' promise. If someone sues us because of something your side did wrong, you agree to pay for our defense and any costs. Think of it like agreeing that whoever spills the milk has to clean it up and pay for a new carton.",

        "What happens if the other party breaches the non-compete agreement":
            "They promised not to start a competing business or work for your rivals. If they break that promise, you can take them to court and ask a judge to make them stop and possibly pay you money for the customers you lost."
    },
    "Medical": {
        "Your MRI shows a benign lesion that we should monitor":
            "We found a small spot on your scan, but it's not cancer—it's harmless. Think of it like a freckle on your skin. We just want to check on it every few months to make sure it stays the same size and doesn't change.",

        "The etiology of your chronic fatigue syndrome":
            "We're trying to find the root cause of why you feel exhausted all the time. It's like being a detective figuring out why a car won't start—is it the battery, the fuel, or something else? Your tiredness could come from a virus you had, stress, or how your immune system is working.",

        "Why we're recommending prophylactic treatment given your comorbidities":
            "Because you have several health conditions at once, we want to give you medicine now to prevent a problem before it happens. It's like putting on sunscreen before going to the beach—we're protecting you ahead of time because you're at higher risk."
    },
    "Financial": {
        "How compound interest and amortization affect your mortgage payments":
            "Your loan grows because you pay interest on interest—money owed on money already owed. Your monthly payment is split between paying down what you borrowed and paying the bank for lending it to you. Early on, most goes to the bank; later, more chips away at what you owe.",

        "Why we recommend diversifying your portfolio with low-liquidity assets":
            "Don't put all your eggs in one basket. We suggest putting some money into things that are harder to sell quickly—like real estate or private businesses—because they often grow more over time, even though you can't turn them into cash overnight like stocks.",

        "The tax implications of depreciation on your rental property":
            "The government lets you pretend your rental building loses value each year on paper, even if it's actually worth more. This 'paper loss' reduces the income you report, so you pay less in taxes now. It's like getting a discount for wear and tear that hasn't really happened yet."
    },
    "Technical/Engineering": {
        "Why our API has high latency and how microservices could help":
            "Our system is slow to respond because everything runs through one big, overloaded program. Imagine one cashier serving an entire grocery store. Breaking it into smaller, specialized services is like opening more checkout lanes—each handles one type of task faster.",

        "The difference between synchronous and asynchronous processing":
            "Synchronous means waiting in line—you can't order your coffee until the person ahead finishes. Asynchronous means you order, step aside, and they call your name when it's ready. The second way lets more people order at once without everyone standing around waiting.",

        "Why we need to refactor the legacy codebase before adding new features":
            "Our old code is like a cluttered garage—you can barely find anything, and adding new stuff just makes the mess worse. We need to organize and clean it up first. Otherwise, every new feature takes twice as long and breaks things we didn't expect."
    }
}

# ============================================================================
# JARGON DICTIONARIES
# ============================================================================

JARGON_DICTIONARIES = {
    "Legal": [
        "liability", "liable", "indemnify", "indemnification", "breach",
        "statute", "damages", "negligence", "herein", "aforementioned",
        "plaintiff", "defendant", "jurisdiction", "arbitration", "tort",
        "fiduciary", "escrow", "lien", "deposition", "stipulation",
        "injunction", "subpoena", "affidavit", "adjudicate", "appellant"
    ],
    "Medical": [
        "prognosis", "diagnosis", "etiology", "pathology", "contraindicated",
        "idiopathic", "nosocomial", "comorbidity", "prophylactic", "benign",
        "malignant", "metastasis", "hemorrhage", "ischemia", "infarction",
        "edema", "necrosis", "lesion", "syndrome", "acute", "chronic",
        "bilateral", "unilateral", "subcutaneous", "intravenous"
    ],
    "Financial": [
        "amortization", "liquidity", "collateral", "derivative", "equity",
        "fiduciary", "hedge", "leverage", "portfolio", "securities",
        "dividend", "depreciation", "liability", "asset", "accrual",
        "arbitrage", "capitalization", "yield", "maturity", "principal",
        "compound", "annuity", "underwriting", "insolvency", "solvency"
    ],
    "Technical/Engineering": [
        "algorithm", "bandwidth", "latency", "throughput", "scalability",
        "deprecated", "refactor", "polymorphism", "encapsulation", "abstraction",
        "iteration", "recursion", "synchronous", "asynchronous", "protocol",
        "middleware", "backend", "frontend", "deployment", "infrastructure",
        "microservices", "containerization", "orchestration", "API", "SDK"
    ]
}

# ============================================================================
# MODELS TO TEST
# ============================================================================

MODELS = {
    "TinyLlama-1.1B": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    "Qwen2-0.5B": "Qwen/Qwen2-0.5B-Instruct",
    "Gemma-2-2B": "google/gemma-2-2b-it",
}

# ============================================================================
# SMC FUNCTIONS
# ============================================================================

def is_safe(text: str, banned_words: list) -> bool:
    text_lower = text.lower()
    for word in banned_words:
        word_lower = word.lower()
        if (f" {word_lower} " in f" {text_lower} " or
            f" {word_lower}." in f" {text_lower}" or
            f" {word_lower}," in f" {text_lower}" or
            f" {word_lower}?" in f" {text_lower}" or
            f" {word_lower}!" in f" {text_lower}" or
            text_lower.startswith(f"{word_lower} ") or
            text_lower.endswith(f" {word_lower}")):
            return False
    return True

def find_jargon_used(text: str, banned_words: list) -> list:
    text_lower = text.lower()
    found = []
    for word in banned_words:
        word_lower = word.lower()
        if (f" {word_lower} " in f" {text_lower} " or
            f" {word_lower}." in f" {text_lower}" or
            f" {word_lower}," in f" {text_lower}" or
            f" {word_lower}?" in f" {text_lower}" or
            f" {word_lower}!" in f" {text_lower}" or
            text_lower.startswith(f"{word_lower} ") or
            text_lower.endswith(f" {word_lower}")):
            found.append(word)
    return list(set(found))

def smc_translate(concept, profession, tokenizer, model, num_particles=5, max_steps=20, tokens_per_step=4):
    banned_words = JARGON_DICTIONARIES.get(profession, [])

    prompt = f"""You are an expert {profession.lower()} professional explaining a concept to a client with no background in your field.

Rules:
- Explain as if talking to a curious 10-year-old
- Use a concrete, relatable real-world example to illustrate the concept
- Avoid redundancy (don't say "X is Y such as Y")
- Keep it concise: 2-3 sentences max

Concept to explain: {concept}

Simple explanation with example:"""

    particles = [prompt]

    for step in range(max_steps):
        candidates = []

        for particle in particles:
            inputs = tokenizer(particle, return_tensors="pt").to(model.device)

            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=tokens_per_step,
                    num_return_sequences=3,
                    do_sample=True,
                    temperature=0.8,
                    top_p=0.9,
                    pad_token_id=tokenizer.eos_token_id
                )

            for out in outputs:
                decoded = tokenizer.decode(out, skip_special_tokens=True)
                candidates.append(decoded)

        valid_candidates = [c for c in candidates if is_safe(c, banned_words)]

        if valid_candidates:
            unique_candidates = list(set(valid_candidates))
            random.shuffle(unique_candidates)
            particles = unique_candidates[:num_particles]
        else:
            break

        current_text = particles[0].split("Simple explanation with example:")[-1].strip()
        # Only stop if we have a good amount of text ending with punctuation
        if current_text.endswith(('.', '!', '?')) and len(current_text) > 100:
            break

    final_text = particles[0].split("Simple explanation with example:")[-1].strip()
    jargon_found = find_jargon_used(final_text, banned_words)

    return final_text, jargon_found

def grade_output(output, benchmark, jargon_found, profession):
    """Grade the output on multiple criteria."""
    scores = {}

    # 1. Jargon-free (0-25 points)
    if len(jargon_found) == 0:
        scores['jargon_free'] = 25
    else:
        scores['jargon_free'] = max(0, 25 - (len(jargon_found) * 10))

    # 2. Has example/analogy (0-25 points)
    example_indicators = ['like', 'imagine', 'think of', 'for example', 'such as', 'similar to', 'as if']
    has_example = any(ind in output.lower() for ind in example_indicators)
    scores['has_example'] = 25 if has_example else 0

    # 3. Appropriate length (0-25 points) - not too short, not too long
    word_count = len(output.split())
    if 20 <= word_count <= 100:
        scores['length'] = 25
    elif 10 <= word_count < 20 or 100 < word_count <= 150:
        scores['length'] = 15
    else:
        scores['length'] = 5

    # 4. Coherence - ends properly (0-25 points)
    if output.strip().endswith(('.', '!', '?')):
        scores['coherence'] = 25
    elif len(output) > 30:
        scores['coherence'] = 15
    else:
        scores['coherence'] = 5

    total = sum(scores.values())
    return scores, total

# ============================================================================
# MAIN TEST RUNNER
# ============================================================================

def run_benchmark_tests(models_to_test=None):
    """Run all benchmark tests and return results."""

    if models_to_test is None:
        models_to_test = list(MODELS.keys())

    results = {
        "timestamp": datetime.now().isoformat(),
        "models": {},
        "summary": {}
    }

    # Load models
    loaded_models = {}
    loaded_tokenizers = {}

    for model_name in models_to_test:
        if model_name not in MODELS:
            print(f"⚠️ Unknown model: {model_name}, skipping...")
            continue

        model_id = MODELS[model_name]
        print(f"\n📦 Loading {model_name}...")

        try:
            loaded_tokenizers[model_name] = AutoTokenizer.from_pretrained(model_id)
            loaded_models[model_name] = AutoModelForCausalLM.from_pretrained(
                model_id,
                device_map="auto",
                torch_dtype=torch.float16
            )
            print(f"✅ {model_name} loaded successfully")
        except Exception as e:
            print(f"❌ Failed to load {model_name}: {e}")
            continue

    # Run tests
    for model_name in loaded_models.keys():
        print(f"\n{'='*60}")
        print(f"🧪 Testing {model_name}")
        print('='*60)

        results["models"][model_name] = {
            "tests": [],
            "total_score": 0,
            "max_possible": 0
        }

        tokenizer = loaded_tokenizers[model_name]
        model = loaded_models[model_name]

        for profession, examples in BENCHMARKS.items():
            for concept, benchmark in examples.items():
                print(f"\n📝 {profession}: {concept[:50]}...")

                try:
                    output, jargon_found = smc_translate(
                        concept, profession, tokenizer, model,
                        num_particles=5, max_steps=25, tokens_per_step=6
                    )

                    scores, total = grade_output(output, benchmark, jargon_found, profession)

                    test_result = {
                        "profession": profession,
                        "concept": concept,
                        "benchmark": benchmark,
                        "output": output,
                        "jargon_found": jargon_found,
                        "scores": scores,
                        "total_score": total
                    }

                    results["models"][model_name]["tests"].append(test_result)
                    results["models"][model_name]["total_score"] += total
                    results["models"][model_name]["max_possible"] += 100

                    print(f"   Output: {output[:100]}...")
                    print(f"   Jargon found: {jargon_found if jargon_found else 'None ✅'}")
                    print(f"   Score: {total}/100")

                except Exception as e:
                    print(f"   ❌ Error: {e}")
                    results["models"][model_name]["tests"].append({
                        "profession": profession,
                        "concept": concept,
                        "error": str(e),
                        "total_score": 0
                    })
                    results["models"][model_name]["max_possible"] += 100

    # Calculate summary
    print(f"\n{'='*60}")
    print("📊 FINAL RESULTS")
    print('='*60)

    for model_name, data in results["models"].items():
        pct = (data["total_score"] / data["max_possible"] * 100) if data["max_possible"] > 0 else 0
        results["summary"][model_name] = {
            "total_score": data["total_score"],
            "max_possible": data["max_possible"],
            "percentage": round(pct, 1)
        }
        print(f"\n{model_name}:")
        print(f"   Total Score: {data['total_score']}/{data['max_possible']} ({pct:.1f}%)")

    # Save results
    with open("benchmark_results.json", "w") as f:
        json.dump(results, f, indent=2)
    print(f"\n💾 Results saved to benchmark_results.json")

    return results

if __name__ == "__main__":
    import sys

    # Allow specifying which models to test via command line
    if len(sys.argv) > 1:
        models_to_test = sys.argv[1:]
    else:
        # Default: test all models
        models_to_test = None

    results = run_benchmark_tests(models_to_test)