better evaluation

Author: adaamko

squeez/training/evaluate.py

@@ -1,50 +1,126 @@
-"""Evaluation script for tool output extraction model.
+"""Evaluation script for squeez tool output extraction model.
 
 Metrics:
-- Line-level Precision/Recall/F1 (vs ground truth relevant lines)
-- ROUGE-L between generated and reference filtered output
-- Compression ratio (output tokens / input tokens)
+- Span Exact Match: fraction of samples where predicted lines == reference lines exactly
+- Span Precision/Recall/F1: line-level set overlap between predicted and reference
+- Empty Accuracy: correctly predicting empty vs non-empty relevant_lines
+- ROUGE-L: token-level overlap between concatenated predicted and reference lines
+- Compression ratio: output lines / input lines
 """
 
 import argparse
 import json
 import logging
-import re
 import statistics
 
 logger = logging.getLogger(__name__)
 
 
-def extract_line_numbers(text: str) -> set[int]:
-    """Extract line numbers from filtered output."""
-    line_nums = set()
-    for line in text.split("\n"):
-        match = re.match(r"^(\d+):", line)
-        if match:
-            line_nums.add(int(match.group(1)))
-    return line_nums
+def _parse_relevant_lines(text: str) -> list[str]:
+    """Parse relevant_lines from model output (JSON or raw text).
+
+    Handles:
+    - Valid JSON: {"relevant_lines": ["line1", "line2"]}
+    - Raw text fallback: split by newlines
+    """
+    text = text.strip()
+    try:
+        data = json.loads(text)
+        lines = data.get("relevant_lines", [])
+        if isinstance(lines, list):
+            return [str(line).strip() for line in lines if str(line).strip()]
+    except (json.JSONDecodeError, TypeError, AttributeError):
+        pass
+
+    # Fallback: treat each non-empty line as a span
+    return [line.strip() for line in text.split("\n") if line.strip()]
 
 
-def compute_line_level_metrics(predicted: str, reference: str) -> dict[str, float]:
-    """Compute line-level precision, recall, and F1."""
-    pred_lines = extract_line_numbers(predicted)
-    ref_lines = extract_line_numbers(reference)
+def compute_span_metrics(predicted: list[str], reference: list[str]) -> dict[str, float]:
+    """Compute span-level precision, recall, F1 using set overlap on normalized lines."""
+    pred_set = set(predicted)
+    ref_set = set(reference)
 
-    if not ref_lines:
-        return {"precision": 0.0, "recall": 0.0, "f1": 0.0}
+    if not ref_set and not pred_set:
+        return {"precision": 1.0, "recall": 1.0, "f1": 1.0, "exact_match": 1.0}
 
-    tp = len(pred_lines & ref_lines)
-    precision = tp / len(pred_lines) if pred_lines else 0.0
-    recall = tp / len(ref_lines) if ref_lines else 0.0
+    if not ref_set or not pred_set:
+        return {"precision": 0.0, "recall": 0.0, "f1": 0.0, "exact_match": 0.0}
+
+    tp = len(pred_set & ref_set)
+    precision = tp / len(pred_set)
+    recall = tp / len(ref_set)
     f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0
+    exact_match = 1.0 if pred_set == ref_set else 0.0
 
     return {
         "precision": round(precision, 4),
         "recall": round(recall, 4),
         "f1": round(f1, 4),
+        "exact_match": exact_match,
     }
 
 
+def compute_partial_overlap(predicted: list[str], reference: list[str]) -> float:
+    """Compute partial overlap ratio using character-level intersection.
+
+    For each reference line, find the best matching predicted line (substring match)
+    and compute the fraction of reference characters covered.
+    """
+    if not reference:
+        return 1.0 if not predicted else 0.0
+    if not predicted:
+        return 0.0
+
+    total_chars = 0
+    matched_chars = 0
+
+    for ref_line in reference:
+        total_chars += len(ref_line)
+        best = 0
+        for pred_line in predicted:
+            # Check substring containment both ways
+            if ref_line in pred_line or pred_line in ref_line:
+                best = max(best, min(len(ref_line), len(pred_line)))
+            else:
+                # Character-level overlap via set intersection on character bigrams
+                ref_bigrams = (
+                    {ref_line[i : i + 2] for i in range(len(ref_line) - 1)}
+                    if len(ref_line) > 1
+                    else {ref_line}
+                )
+                pred_bigrams = (
+                    {pred_line[i : i + 2] for i in range(len(pred_line) - 1)}
+                    if len(pred_line) > 1
+                    else {pred_line}
+                )
+                if ref_bigrams:
+                    overlap = len(ref_bigrams & pred_bigrams) / len(ref_bigrams)
+                    best = max(best, int(overlap * len(ref_line)))
+        matched_chars += best
+
+    return round(matched_chars / total_chars, 4) if total_chars > 0 else 0.0
+
+
+def compute_empty_accuracy(predicted: list[str], reference: list[str]) -> dict[str, float | str]:
+    """Check if model correctly predicts empty vs non-empty.
+
+    Returns category (true_positive, true_negative, false_positive, false_negative)
+    and whether correct.
+    """
+    ref_empty = len(reference) == 0
+    pred_empty = len(predicted) == 0
+
+    if ref_empty and pred_empty:
+        return {"category": "true_negative", "correct": 1.0}
+    elif ref_empty and not pred_empty:
+        return {"category": "false_positive", "correct": 0.0}
+    elif not ref_empty and pred_empty:
+        return {"category": "false_negative", "correct": 0.0}
+    else:
+        return {"category": "true_positive", "correct": 1.0}
+
+
 def _lcs_length(x: list[str], y: list[str]) -> int:
     """Compute longest common subsequence length."""
     m, n = len(x), len(y)
@@ -67,8 +143,8 @@ def compute_rouge_l(predicted: str, reference: str) -> float:
         return 0.0
 
     lcs = _lcs_length(pred_tokens, ref_tokens)
-    precision = lcs / len(pred_tokens) if pred_tokens else 0.0
-    recall = lcs / len(ref_tokens) if ref_tokens else 0.0
+    precision = lcs / len(pred_tokens)
+    recall = lcs / len(ref_tokens)
 
     if precision + recall == 0:
         return 0.0
@@ -100,7 +176,6 @@ def evaluate_model(
 
     Returns:
         Dict with aggregate metrics
-
     """
     import torch
     from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -115,6 +190,9 @@ def evaluate_model(
     )
     model.eval()
 
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
     # Load eval data
     samples = []
     with open(eval_file) as f:
@@ -126,19 +204,29 @@ def evaluate_model(
     logger.info(f"Evaluating on {len(samples)} samples")
 
     all_metrics = {
-        "line_precision": [],
-        "line_recall": [],
-        "line_f1": [],
+        "span_precision": [],
+        "span_recall": [],
+        "span_f1": [],
+        "exact_match": [],
+        "partial_overlap": [],
+        "empty_accuracy": [],
         "rouge_l": [],
         "compression": [],
     }
 
+    empty_confusion = {
+        "true_positive": 0,
+        "true_negative": 0,
+        "false_positive": 0,
+        "false_negative": 0,
+    }
+
     for i, sample in enumerate(samples):
         prompt = sample["prompt"]
-        reference = sample["response"]
+        reference_raw = sample["response"]
 
         # Generate
-        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=16384)
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
 
         with torch.no_grad():
@@ -150,26 +238,47 @@ def evaluate_model(
                 pad_token_id=tokenizer.pad_token_id,
             )
 
-        generated = tokenizer.decode(
+        generated_raw = tokenizer.decode(
             outputs[0][inputs["input_ids"].shape[1] :],
             skip_special_tokens=True,
         )
 
-        # Compute metrics
-        line_metrics = compute_line_level_metrics(generated, reference)
-        rouge = compute_rouge_l(generated, reference)
-        original_output = prompt.split("<|user|>")[-1].split("<|assistant|>")[0]
-        compression = compute_compression_ratio(original_output, generated)
-
-        all_metrics["line_precision"].append(line_metrics["precision"])
-        all_metrics["line_recall"].append(line_metrics["recall"])
-        all_metrics["line_f1"].append(line_metrics["f1"])
+        # Parse both into line lists
+        pred_lines = _parse_relevant_lines(generated_raw)
+        ref_lines = _parse_relevant_lines(reference_raw)
+
+        # Span metrics
+        span = compute_span_metrics(pred_lines, ref_lines)
+        all_metrics["span_precision"].append(span["precision"])
+        all_metrics["span_recall"].append(span["recall"])
+        all_metrics["span_f1"].append(span["f1"])
+        all_metrics["exact_match"].append(span["exact_match"])
+
+        # Partial overlap
+        partial = compute_partial_overlap(pred_lines, ref_lines)
+        all_metrics["partial_overlap"].append(partial)
+
+        # Empty accuracy
+        empty = compute_empty_accuracy(pred_lines, ref_lines)
+        all_metrics["empty_accuracy"].append(empty["correct"])
+        empty_confusion[empty["category"]] += 1
+
+        # ROUGE-L on concatenated text
+        pred_text = "\n".join(pred_lines)
+        ref_text = "\n".join(ref_lines)
+        rouge = compute_rouge_l(pred_text, ref_text)
         all_metrics["rouge_l"].append(rouge)
+
+        # Compression
+        original_output = prompt.split("<|user|>")[-1].split("<|assistant|>")[0]
+        compression = compute_compression_ratio(original_output, pred_text)
         all_metrics["compression"].append(compression)
 
         if (i + 1) % 10 == 0:
             logger.info(
-                f"  [{i + 1}/{len(samples)}] F1={line_metrics['f1']:.3f} ROUGE-L={rouge:.3f}"
+                f"  [{i + 1}/{len(samples)}] "
+                f"F1={span['f1']:.3f} EM={span['exact_match']:.0f} "
+                f"ROUGE-L={rouge:.3f}"
             )
 
     # Aggregate
@@ -182,11 +291,16 @@ def evaluate_model(
                 "stdev": round(statistics.stdev(values), 4) if len(values) > 1 else 0,
             }
 
-    logger.info("=" * 50)
+    results["empty_confusion"] = empty_confusion
+    results["num_samples"] = len(samples)
+
+    logger.info("=" * 60)
     logger.info("EVALUATION RESULTS")
-    logger.info("=" * 50)
+    logger.info("=" * 60)
     for key, stats in results.items():
-        logger.info(f"  {key}: mean={stats['mean']:.4f} median={stats['median']:.4f}")
+        if isinstance(stats, dict) and "mean" in stats:
+            logger.info(f"  {key:20s}: mean={stats['mean']:.4f}  median={stats['median']:.4f}")
+    logger.info(f"  {'empty_confusion':20s}: {empty_confusion}")
 
     return results
 
@@ -203,7 +317,7 @@ def build_parser(parser: argparse.ArgumentParser | None = None) -> argparse.Argu
         required=True,
         help="Path to the trained extractor model",
     )
-    parser.add_argument("--eval-file", required=True, help="Path to eval.jsonl")
+    parser.add_argument("--eval-file", required=True, help="Path to test.jsonl")
     parser.add_argument("--max-samples", type=int, default=None)
     parser.add_argument("--max-new-tokens", type=int, default=1024)
     return parser