eval_aes_accuracy.py

#!/usr/bin/env python3
"""
Evaluate the AES LoRA model on the validation set.
Computes exact-match accuracy and token-level similarity.

Can evaluate either:
  1. Merged model (no LoRA adapter needed) — set --merged-path
  2. Base + LoRA adapter — set --base-path and --adapter-path

Usage:
  python3 eval_aes_accuracy.py
  python3 eval_aes_accuracy.py --base-path /path/to/merged --adapter-path /path/to/aes_lora
  python3 eval_aes_accuracy.py --merged-path /path/to/merged_aes_final
  python3 eval_aes_accuracy.py --max-samples 50
"""
import argparse
import json
import re
import time
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

WORKSPACE = Path("/workspace/elinnos")
DEFAULT_VAL_DATA = WORKSPACE / "aes_training" / "data" / "aes_val.jsonl"
DEFAULT_BASE = WORKSPACE / "merged_models" / "elinnos_all_merged_final"
DEFAULT_ADAPTER = WORKSPACE / "elinnos-qwen2.5-7b-aes-lora"
CHAT_TEMPLATE_SRC = WORKSPACE / "elinnos-qwen2.5-7b-multi-ip-lora-v4" / "chat_template.jinja"

MAX_NEW_TOKENS = 8192
TEMPERATURE = 0.2

_DOLLAR_TAG_RE = re.compile(r'([A-Za-z_][A-Za-z0-9_]*)\$\$([A-Za-z0-9]+)')


def normalize_dollar_tags(text: str) -> str:
    first_tag: dict[str, str] = {}
    def _repl(m: "re.Match[str]") -> str:
        base, tag = m.group(1), m.group(2)
        canonical = first_tag.setdefault(base, tag)
        return f"{base}$${canonical}"
    return _DOLLAR_TAG_RE.sub(_repl, text)


def strip_dollar_tags(text: str) -> str:
    return _DOLLAR_TAG_RE.sub(r"\1", text)


def normalize_whitespace(text: str) -> str:
    return re.sub(r'\s+', ' ', text).strip()


def compute_accuracy(expected: str, generated: str) -> dict:
    import difflib
    exp_norm = normalize_whitespace(strip_dollar_tags(expected))
    gen_norm = normalize_whitespace(strip_dollar_tags(generated))

    exact_match = (exp_norm == gen_norm)

    exp_words = exp_norm.split()
    gen_words = gen_norm.split()
    if len(exp_words) == 0:
        token_acc = 1.0 if len(gen_words) == 0 else 0.0
    else:
        matcher = difflib.SequenceMatcher(None, exp_words, gen_words)
        token_acc = matcher.ratio()

    char_sim = difflib.SequenceMatcher(None, exp_norm, gen_norm).ratio()

    return {
        "exact_match": exact_match,
        "token_similarity": token_acc,
        "char_similarity": char_sim,
        "exp_len": len(exp_words),
        "gen_len": len(gen_words),
    }


def parse_args():
    p = argparse.ArgumentParser(description="Evaluate AES LoRA model accuracy")
    p.add_argument("--val-data", type=str, default=str(DEFAULT_VAL_DATA))
    p.add_argument("--base-path", type=str, default=str(DEFAULT_BASE))
    p.add_argument("--adapter-path", type=str, default=str(DEFAULT_ADAPTER))
    p.add_argument("--merged-path", type=str, default=None,
                   help="If set, load this merged model directly (no adapter)")
    p.add_argument("--max-samples", type=int, default=None,
                   help="Limit number of samples for quick eval")
    p.add_argument("--max-new-tokens", type=int, default=MAX_NEW_TOKENS)
    p.add_argument("--temperature", type=float, default=TEMPERATURE)
    p.add_argument("--output", type=str, default=str(WORKSPACE / "elinnos-qwen2.5-7b-aes-lora" / "eval_results.json"))
    return p.parse_args()


def main():
    args = parse_args()
    import difflib

    val_data = Path(args.val_data).resolve()
    output_path = Path(args.output).resolve()

    print("=" * 70)
    print("  AES SECURITY IP — ACCURACY EVALUATION")
    print("=" * 70)

    # Load validation samples
    val_samples = []
    with val_data.open() as f:
        for line in f:
            line = line.strip()
            if line:
                val_samples.append(json.loads(line))

    if args.max_samples:
        val_samples = val_samples[:args.max_samples]

    print(f"  Val data: {val_data}")
    print(f"  Samples:  {len(val_samples)}")
    print(f"  Temperature: {args.temperature}")
    print(f"  Max new tokens: {args.max_new_tokens}")

    # Load model
    if args.merged_path:
        model_path = Path(args.merged_path).resolve()
        print(f"  Model: {model_path} (merged, no adapter)")
        tokenizer = AutoTokenizer.from_pretrained(str(model_path), trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        ct = model_path / "chat_template.jinja"
        if not ct.exists():
            ct = CHAT_TEMPLATE_SRC
        if ct.exists():
            tokenizer.chat_template = ct.read_text()
            print(f"  Chat template: {ct}")

        print("Loading merged model (bf16)...")
        model = AutoModelForCausalLM.from_pretrained(
            str(model_path), torch_dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True,
        )
    else:
        base_path = Path(args.base_path).resolve()
        adapter_path = Path(args.adapter_path).resolve()
        print(f"  Base: {base_path}")
        print(f"  Adapter: {adapter_path}")

        tokenizer = AutoTokenizer.from_pretrained(str(base_path), trust_remote_code=True)
        if tokenizer.pad_token is None:
            tokenizer.pad_token = tokenizer.eos_token
        ct = adapter_path / "chat_template.jinja"
        if not ct.exists():
            ct = base_path / "chat_template.jinja"
        if not ct.exists():
            ct = CHAT_TEMPLATE_SRC
        if ct.exists():
            tokenizer.chat_template = ct.read_text()
            print(f"  Chat template: {ct}")

        print("Loading base model (bf16)...")
        model = AutoModelForCausalLM.from_pretrained(
            str(base_path), torch_dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True, low_cpu_mem_usage=True,
        )
        print(f"Applying LoRA adapter from {adapter_path}...")
        model = PeftModel.from_pretrained(model, str(adapter_path))

    model.eval()
    if torch.cuda.is_available():
        print(f"  GPU: {torch.cuda.get_device_name(0)}")
    print()

    results = []
    exact_matches = 0
    total_token_sim = 0.0
    total_char_sim = 0.0

    for i, sample in enumerate(val_samples):
        messages = sample["messages"]
        system_msg = next(m["content"] for m in messages if m["role"] == "system")
        user_msg = next(m["content"] for m in messages if m["role"] == "user")
        expected = next(m["content"] for m in messages if m["role"] == "assistant")

        prompt_messages = [
            {"role": "system", "content": system_msg},
            {"role": "user", "content": user_msg},
        ]

        text = tokenizer.apply_chat_template(prompt_messages, tokenize=False, add_generation_prompt=True)
        inputs = tokenizer(text, return_tensors="pt").to(model.device)

        start = time.time()
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=args.max_new_tokens,
                temperature=args.temperature,
                do_sample=args.temperature > 0,
                top_p=0.9,
                repetition_penalty=1.05,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        elapsed = time.time() - start

        input_len = inputs["input_ids"].shape[1]
        generated = tokenizer.decode(outputs[0][input_len:], skip_special_tokens=True)
        generated = normalize_dollar_tags(generated)

        metrics = compute_accuracy(expected, generated)
        results.append({
            "idx": i,
            "user_prompt": user_msg[:100],
            "exact_match": metrics["exact_match"],
            "token_similarity": metrics["token_similarity"],
            "char_similarity": metrics["char_similarity"],
            "exp_len": metrics["exp_len"],
            "gen_len": metrics["gen_len"],
            "gen_time": elapsed,
        })

        if metrics["exact_match"]:
            exact_matches += 1
        total_token_sim += metrics["token_similarity"]
        total_char_sim += metrics["char_similarity"]

        status = "EXACT" if metrics["exact_match"] else f"tok_sim={metrics['token_similarity']:.3f}"
        print(f"  [{i+1}/{len(val_samples)}] {status}  ({elapsed:.1f}s, exp={metrics['exp_len']}w gen={metrics['gen_len']}w)")

    n = len(val_samples)
    print("\n" + "=" * 70)
    print("  ACCURACY RESULTS")
    print("=" * 70)
    print(f"  Total samples:       {n}")
    print(f"  Exact matches:       {exact_matches}/{n} ({exact_matches/n*100:.1f}%)")
    print(f"  Avg token similarity: {total_token_sim/n*100:.1f}%")
    print(f"  Avg char similarity:  {total_char_sim/n*100:.1f}%")
    print("=" * 70)

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("w") as f:
        json.dump({
            "summary": {
                "total_samples": n,
                "exact_matches": exact_matches,
                "exact_match_rate": exact_matches / n,
                "avg_token_similarity": total_token_sim / n,
                "avg_char_similarity": total_char_sim / n,
            },
            "details": results,
        }, f, indent=2)
    print(f"\nDetailed results saved to: {output_path}")


if __name__ == "__main__":
    main()