evaluate.py

#!/usr/bin/env python3
"""
Evaluate the fine-tuned model on the test set.

Generates a commit message for each test diff and compares to the reference.
Prints ROUGE-L scores and logs examples for manual review.
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

from build_dataset import USER_PREFIX, USER_SUFFIX

DEFAULT_MODEL = "Qwen/Qwen3-0.6B"
DEFAULT_ADAPTER = "adapters"


def _rouge_l(reference: str, hypothesis: str) -> float:
    """Compute ROUGE-L F1 between two strings (word-level LCS)."""
    ref_words = reference.lower().split()
    hyp_words = hypothesis.lower().split()
    if not ref_words or not hyp_words:
        return 0.0
    # LCS via DP
    m, n = len(ref_words), len(hyp_words)
    dp = [[0] * (n + 1) for _ in range(m + 1)]
    for i in range(1, m + 1):
        for j in range(1, n + 1):
            if ref_words[i - 1] == hyp_words[j - 1]:
                dp[i][j] = dp[i - 1][j - 1] + 1
            else:
                dp[i][j] = max(dp[i - 1][j], dp[i][j - 1])
    lcs_len = dp[m][n]
    precision = lcs_len / n
    recall = lcs_len / m
    if precision + recall == 0:
        return 0.0
    return 2 * precision * recall / (precision + recall)


def main() -> int:
    parser = argparse.ArgumentParser(description="Evaluate model on test set.")
    parser.add_argument(
        "--test-file", type=Path, default=Path("data/test.jsonl"),
        help="Test JSONL file (default: data/test.jsonl)",
    )
    parser.add_argument(
        "--model", default=DEFAULT_MODEL,
        help=f"Base model (default: {DEFAULT_MODEL})",
    )
    parser.add_argument(
        "--adapter-path", default=DEFAULT_ADAPTER,
        help=f"LoRA adapter directory (default: {DEFAULT_ADAPTER})",
    )
    parser.add_argument(
        "--no-adapter", action="store_true",
        help="Evaluate base model without adapter (baseline comparison)",
    )
    parser.add_argument(
        "--temp", type=float, default=0.0,
        help="Sampling temperature (default: 0.0 = greedy)",
    )
    parser.add_argument(
        "--max-tokens", type=int, default=100,
        help="Max tokens to generate (default: 100)",
    )
    parser.add_argument(
        "--output", type=Path, default=None,
        help="Write detailed results as JSONL to this file",
    )
    args = parser.parse_args()

    if not args.test_file.is_file():
        print(f"error: test file not found: {args.test_file}", file=sys.stderr)
        return 1

    examples = []
    with args.test_file.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
                examples.append(json.loads(line))

    if not examples:
        print("error: no test examples found", file=sys.stderr)
        return 1

    print(f"Evaluating on {len(examples)} test examples...")

    from mlx_lm import load, generate

    adapter = None if args.no_adapter else args.adapter_path
    model, tokenizer = load(args.model, adapter_path=adapter)

    scores = []
    results = []
    for i, ex in enumerate(examples):
        diff = ex["diff"]
        reference = ex["message"]
        prompt = USER_PREFIX + diff + USER_SUFFIX

        response = generate(
            model, tokenizer, prompt=prompt,
            max_tokens=args.max_tokens, temp=args.temp,
        )
        hypothesis = response.split("<|im_end|>")[0].split("<|endoftext|>")[0].strip()

        score = _rouge_l(reference, hypothesis)
        scores.append(score)

        result = {
            "reference": reference,
            "generated": hypothesis,
            "rouge_l": round(score, 4),
        }
        results.append(result)

        tag = "GOOD" if score >= 0.5 else "OK  " if score >= 0.2 else "MISS"
        print(f"  [{i+1:>3}/{len(examples)}] {tag}  ROUGE-L={score:.3f}")
        print(f"           ref: {reference}")
        print(f"           gen: {hypothesis}")
        print()

    avg = sum(scores) / len(scores) if scores else 0
    good = sum(1 for s in scores if s >= 0.5)
    ok = sum(1 for s in scores if 0.2 <= s < 0.5)
    miss = sum(1 for s in scores if s < 0.2)

    print("=" * 60)
    print(f"Results ({len(scores)} examples)")
    print("=" * 60)
    print(f"Mean ROUGE-L:  {avg:.4f}")
    print(f"GOOD (≥0.5):   {good}/{len(scores)} ({100*good/len(scores):.0f}%)")
    print(f"OK   (≥0.2):   {ok}/{len(scores)} ({100*ok/len(scores):.0f}%)")
    print(f"MISS (<0.2):   {miss}/{len(scores)} ({100*miss/len(scores):.0f}%)")

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        with args.output.open("w", encoding="utf-8") as f:
            for r in results:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"\nDetailed results: {args.output}")

    return 0


if __name__ == "__main__":
    raise SystemExit(main())