05_evaluate.py

"""
Step 05: Evaluate base vs fine-tuned model on MISRA compliance.

Generates code from both the base Qwen model and our fine-tuned LoRA adapter,
then runs cppcheck MISRA on each output to compare violation counts.

Inference runs on Modal (GPU). cppcheck runs locally.

Requires: Modal authentication, cppcheck installed locally.

Usage:
    .venv/bin/modal run 05_evaluate.py
"""
import json
import os
import re

import modal
import yaml

app = modal.App("misra-coder-eval")

CHECKPOINT = "checkpoint-230"

SYSTEM_PROMPT = (
    "You are a C++ developer that writes strictly MISRA C:2012 compliant code. "
    "All code you produce must pass cppcheck MISRA analysis with zero violations."
)

eval_image = (
    modal.Image.debian_slim(python_version="3.11")
    .uv_pip_install(
        "accelerate==1.9.0",
        "datasets==3.6.0",
        "hf-transfer==0.1.9",
        "huggingface_hub==0.34.2",
        "peft==0.16.0",
        "transformers==4.54.0",
        "trl==0.19.1",
        "unsloth[cu128-torch270]==2025.7.8",
        "unsloth_zoo==2025.7.10",
    )
    .env({"HF_HOME": "/model_cache"})
)

with eval_image.imports():
    import unsloth  # noqa: F401,I001
    import torch
    from unsloth import FastLanguageModel

model_cache = modal.Volume.from_name("misra-model-cache", create_if_missing=True)
checkpoint_vol = modal.Volume.from_name("misra-checkpoints", create_if_missing=True)


@app.function(
    image=eval_image,
    gpu="L40S",
    volumes={
        "/model_cache": model_cache,
        "/checkpoints": checkpoint_vol,
    },
    timeout=1 * 60 * 60,
)
def generate_code(prompts, max_tokens, lora_path=None):
    if lora_path:
        print(f"Loading fine-tuned model from {lora_path}...")
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=lora_path,
            max_seq_length=2048,
            dtype=None,
            load_in_4bit=True,
        )
    else:
        print("Loading base Qwen2.5-Coder-7B-Instruct...")
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name="Qwen/Qwen2.5-Coder-7B-Instruct",
            max_seq_length=2048,
            dtype=None,
            load_in_4bit=True,
        )

    FastLanguageModel.for_inference(model)

    outputs = []
    for i, prompt in enumerate(prompts):
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt},
        ]
        inputs = tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        ).to("cuda")

        with torch.no_grad():
            output_ids = model.generate(
                input_ids=inputs,
                max_new_tokens=max_tokens,
                temperature=0.2,
                do_sample=True,
            )

        response = tokenizer.decode(
            output_ids[0][inputs.shape[1]:], skip_special_tokens=True
        )
        outputs.append(response)

        if (i + 1) % 10 == 0:
            print(f"  Generated {i + 1}/{len(prompts)}")

    print(f"  Done: {len(outputs)} outputs")
    return outputs


def extract_code(text):
    match = re.search(r"```(?:cpp|c\+\+|c)?\s*\n(.*?)```", text, re.DOTALL)
    if match:
        return match.group(1).strip()
    return text.strip()


@app.local_entrypoint()
def main():
    with open("config.yaml") as f:
        cfg = yaml.safe_load(f)

    eval_cfg = cfg["eval"]
    training_dir = cfg["paths"]["training_data"]

    prompts = []
    with open(os.path.join(training_dir, "val.jsonl")) as f:
        for line in f:
            data = json.loads(line)
            for msg in data["messages"]:
                if msg["role"] == "user":
                    prompts.append(msg["content"])
                    break

    num_prompts = min(eval_cfg["num_prompts"], len(prompts))
    prompts = prompts[:num_prompts]

    print(f"Evaluating {num_prompts} prompts")
    print(f"Checkpoint: {CHECKPOINT}\n")

    print("Generating with base model...")
    base_outputs = generate_code.remote(prompts, eval_cfg["max_tokens"])

    lora_path = f"/checkpoints/misra-coder/{CHECKPOINT}"
    print(f"Generating with fine-tuned model...")
    ft_outputs = generate_code.remote(prompts, eval_cfg["max_tokens"], lora_path=lora_path)

    from utils.cppcheck_runner import run_misra_check

    print("Running cppcheck MISRA analysis...\n")

    results = []
    for i, (prompt, base_raw, ft_raw) in enumerate(zip(prompts, base_outputs, ft_outputs)):
        base_code = extract_code(base_raw)
        ft_code = extract_code(ft_raw)

        base_v = run_misra_check(base_code)
        ft_v = run_misra_check(ft_code)

        base_count = len(base_v) if base_v else 0
        ft_count = len(ft_v) if ft_v else 0

        results.append({
            "prompt": prompt,
            "base_violations": base_count,
            "ft_violations": ft_count,
            "base_code": base_code,
            "ft_code": ft_code,
        })

        if ft_count < base_count:
            mark = "+"
        elif ft_count == base_count:
            mark = "="
        else:
            mark = "-"

        print(f"  [{i+1:3d}/{num_prompts}] Base: {base_count:2d} | FT: {ft_count:2d} | {mark} {prompt[:55]}")

    base_total = sum(r["base_violations"] for r in results)
    ft_total = sum(r["ft_violations"] for r in results)
    base_avg = base_total / len(results) if results else 0
    ft_avg = ft_total / len(results) if results else 0
    base_zero = sum(1 for r in results if r["base_violations"] == 0)
    ft_zero = sum(1 for r in results if r["ft_violations"] == 0)
    improved = sum(1 for r in results if r["ft_violations"] < r["base_violations"])
    same = sum(1 for r in results if r["ft_violations"] == r["base_violations"])
    worse = sum(1 for r in results if r["ft_violations"] > r["base_violations"])

    print(f"\n{'='*60}")
    print(f"EVALUATION RESULTS ({num_prompts} prompts)")
    print(f"{'='*60}")
    print(f"                    Base Model    Fine-tuned")
    print(f"  Total violations: {base_total:>10d}    {ft_total:>10d}")
    print(f"  Avg per file:     {base_avg:>10.1f}    {ft_avg:>10.1f}")
    print(f"  Zero-violation:   {base_zero:>10d}    {ft_zero:>10d}")
    print(f"")
    print(f"  Improved: {improved}/{num_prompts}")
    print(f"  Same:     {same}/{num_prompts}")
    print(f"  Worse:    {worse}/{num_prompts}")
    if base_total > 0:
        reduction = (1 - ft_total / base_total) * 100
        print(f"\n  Violation reduction: {reduction:.0f}%")

    output_dir = eval_cfg["output_dir"]
    os.makedirs(output_dir, exist_ok=True)
    results_path = os.path.join(output_dir, "eval_results.json")
    with open(results_path, "w") as f:
        json.dump({
            "summary": {
                "num_prompts": num_prompts,
                "checkpoint": CHECKPOINT,
                "base_total_violations": base_total,
                "ft_total_violations": ft_total,
                "base_avg_violations": base_avg,
                "ft_avg_violations": ft_avg,
                "base_zero_violation_files": base_zero,
                "ft_zero_violation_files": ft_zero,
                "improved": improved,
                "same": same,
                "worse": worse,
            },
            "results": results,
        }, f, indent=2)

    print(f"\n  Full results saved to {results_path}")