Speed up and resume OOD evaluation with batched dynamic generation

scripts/evaluate_model.py

@@ -1,16 +1,18 @@
 #!/usr/bin/env python3
-"""Generation evaluation for TMF921 intent-to-config models/adapters.
+"""Fast generation evaluation for TMF921 intent-to-config models/adapters.
 
 Metrics: JSON parse rate, exact canonical JSON match, field-level F1, simple metadata constraints,
 stratified by split/target_layer/slice_type/lifecycle_operation.
+
+The evaluator is resumable and batched. It periodically writes partial predictions so a long OOD
+run can be stopped/restarted without losing completed examples.
 """
 import argparse
-import csv
 import json
 import os
 from collections import defaultdict
 from pathlib import Path
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
 
 import torch
 from datasets import load_dataset
@@ -29,7 +31,12 @@ def parse_args():
     p.add_argument("--splits", nargs="+", default=["test_in_distribution", "test_template_ood", "test_use_case_ood", "test_sector_ood", "test_adversarial"])
     p.add_argument("--output_dir", default="eval_outputs")
     p.add_argument("--max_samples_per_split", type=int, default=None)
-    p.add_argument("--max_new_tokens", type=int, default=2048)
+    p.add_argument("--max_new_tokens", type=int, default=1536, help="Hard cap for generation. 1536 covers audited dataset outputs with margin.")
+    p.add_argument("--gold_length_buffer", type=int, default=96, help="Dynamic cap = max gold output tokens in batch + buffer, clipped by max_new_tokens")
+    p.add_argument("--batch_size", type=int, default=4, help="Batched generation size. Use 1 if OOM; 4 is usually safe on RTX 6000 Ada for Qwen3-8B QLoRA.")
+    p.add_argument("--save_every", type=int, default=25, help="Write partial predictions every N completed examples")
+    p.add_argument("--resume", action="store_true", default=True, help="Resume from existing predictions.json if present")
+    p.add_argument("--no_resume", dest="resume", action="store_false")
     p.add_argument("--temperature", type=float, default=0.0)
     p.add_argument("--top_p", type=float, default=1.0)
     p.add_argument("--load_in_4bit", action="store_true", help="Load base model in 4-bit for adapter evaluation")
@@ -39,25 +46,43 @@ def parse_args():
 
 
 def make_prompt_messages(messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
-    # Keep all messages before the final assistant answer. For standard rows this is system+user;
-    # for multi-turn rows this preserves earlier assistant turns but removes gold final answer.
     out = []
-    for m in messages:
-        if m.get("role") == "assistant" and m == messages[-1]:
+    for i, m in enumerate(messages):
+        if i == len(messages) - 1 and m.get("role") == "assistant":
             break
         out.append({"role": m.get("role"), "content": m.get("content", "")})
-    # Fallback for any unexpected format.
     if not out:
         out = [m for m in messages if m.get("role") != "assistant"]
     return out
 
 
-def generate_one(model, tokenizer, messages, args):
-    prompt_messages = make_prompt_messages(messages)
-    text = tokenizer.apply_chat_template(prompt_messages, tokenize=False, add_generation_prompt=True)
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
+def make_prompt_text(tokenizer, messages: List[Dict[str, str]]) -> str:
+    return tokenizer.apply_chat_template(make_prompt_messages(messages), tokenize=False, add_generation_prompt=True)
+
+
+def gold_text(example: Dict[str, Any]) -> str:
+    return example.get("completion") or get_message(example["messages"], "assistant")
+
+
+def dynamic_max_new_tokens(tokenizer, examples: List[Dict[str, Any]], args) -> int:
+    lens = []
+    for ex in examples:
+        ids = tokenizer(gold_text(ex), add_special_tokens=False)["input_ids"]
+        lens.append(len(ids))
+    return max(16, min(int(args.max_new_tokens), max(lens) + int(args.gold_length_buffer)))
+
+
+def generate_batch(model, tokenizer, examples: List[Dict[str, Any]], args) -> List[str]:
+    texts = [make_prompt_text(tokenizer, ex["messages"]) for ex in examples]
+    old_padding_side = tokenizer.padding_side
+    tokenizer.padding_side = "left"
+    try:
+        inputs = tokenizer(texts, return_tensors="pt", padding=True).to(model.device)
+    finally:
+        tokenizer.padding_side = old_padding_side
+    max_new = dynamic_max_new_tokens(tokenizer, examples, args)
     gen_kwargs = dict(
-        max_new_tokens=args.max_new_tokens,
+        max_new_tokens=max_new,
         do_sample=args.temperature > 0,
         temperature=args.temperature if args.temperature > 0 else None,
         top_p=args.top_p,
@@ -65,16 +90,16 @@ def generate_one(model, tokenizer, messages, args):
         eos_token_id=tokenizer.eos_token_id,
     )
     gen_kwargs = {k: v for k, v in gen_kwargs.items() if v is not None}
-    with torch.no_grad():
+    with torch.inference_mode():
         out = model.generate(**inputs, **gen_kwargs)
-    new_tokens = out[0, inputs["input_ids"].shape[1]:]
-    return tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
+    new_tokens = out[:, inputs["input_ids"].shape[1]:]
+    return tokenizer.batch_decode(new_tokens, skip_special_tokens=True)
 
 
 def row_metrics(example: Dict[str, Any], prediction: str) -> Dict[str, Any]:
-    gold_text = example.get("completion") or get_message(example["messages"], "assistant")
+    gold = gold_text(example)
     pred_obj, pred_err = parse_json(prediction)
-    gold_obj, gold_err = parse_json(gold_text)
+    gold_obj, gold_err = parse_json(gold)
     out: Dict[str, Any] = {
         "id": example.get("id"),
         "target_layer": example.get("target_layer"),
@@ -84,7 +109,7 @@ def row_metrics(example: Dict[str, Any], prediction: str) -> Dict[str, Any]:
         "gold_parse_json": gold_obj is not None,
         "exact_match": False,
         "prediction": prediction,
-        "gold": gold_text,
+        "gold": gold,
         "parse_error": pred_err,
     }
     if pred_obj is not None and gold_obj is not None:
@@ -108,6 +133,26 @@ def require_cuda():
     print(f"cuda device_count={torch.cuda.device_count()} gpu0={torch.cuda.get_device_name(0)}")
 
 
+def load_existing_predictions(path: Path) -> Tuple[List[Dict[str, Any]], set]:
+    if path.exists():
+        rows = json.loads(path.read_text())
+        done = {str(r.get("id")) for r in rows}
+        return rows, done
+    return [], set()
+
+
+def write_split_outputs(split_dir: Path, rows: List[Dict[str, Any]]) -> Dict[str, Any]:
+    write_json(split_dir / "predictions.json", rows)
+    summary = aggregate_metrics(rows)
+    for key in ["target_layer", "slice_type", "lifecycle_operation"]:
+        groups = defaultdict(list)
+        for r in rows:
+            groups[str(r.get(key))].append(r)
+        summary[f"by_{key}"] = {g: aggregate_metrics(v) for g, v in sorted(groups.items())}
+    write_json(split_dir / "metrics.json", summary)
+    return summary
+
+
 def main():
     args = parse_args()
     require_cuda()
@@ -143,26 +188,36 @@ def main():
         split_ds = ds[split]
         if args.max_samples_per_split:
             split_ds = split_ds.select(range(min(args.max_samples_per_split, len(split_ds))))
-        rows = []
-        print(f"Evaluating {split}: {len(split_ds)} examples")
-        for ex in tqdm(split_ds, desc=split):
-            pred = generate_one(model, tokenizer, ex["messages"], args)
-            rows.append(row_metrics(ex, pred))
-
         split_dir = out_dir / split
         split_dir.mkdir(parents=True, exist_ok=True)
-        write_json(split_dir / "predictions.json", rows)
-        summary = aggregate_metrics(rows)
-        # Stratified summaries.
-        for key in ["target_layer", "slice_type", "lifecycle_operation"]:
-            groups = defaultdict(list)
-            for r in rows:
-                groups[str(r.get(key))].append(r)
-            summary[f"by_{key}"] = {g: aggregate_metrics(v) for g, v in sorted(groups.items())}
-        write_json(split_dir / "metrics.json", summary)
+        pred_path = split_dir / "predictions.json"
+        rows, done_ids = load_existing_predictions(pred_path) if args.resume else ([], set())
+        todo = [ex for ex in split_ds if str(ex.get("id")) not in done_ids]
+        print(f"Evaluating {split}: total={len(split_ds)} already_done={len(done_ids)} remaining={len(todo)} batch_size={args.batch_size}")
+        pbar = tqdm(total=len(todo), desc=split)
+        completed_since_save = 0
+        for start in range(0, len(todo), args.batch_size):
+            batch = todo[start:start + args.batch_size]
+            try:
+                preds = generate_batch(model, tokenizer, batch, args)
+            except torch.cuda.OutOfMemoryError:
+                torch.cuda.empty_cache()
+                if args.batch_size == 1 or len(batch) == 1:
+                    raise
+                preds = []
+                for ex in batch:
+                    preds.extend(generate_batch(model, tokenizer, [ex], args))
+            for ex, pred in zip(batch, preds):
+                rows.append(row_metrics(ex, pred.strip()))
+            pbar.update(len(batch))
+            completed_since_save += len(batch)
+            if completed_since_save >= args.save_every:
+                write_split_outputs(split_dir, rows)
+                completed_since_save = 0
+        pbar.close()
+        summary = write_split_outputs(split_dir, rows)
         all_summary[split] = summary
-
-    write_json(out_dir / "all_metrics.json", all_summary)
+        write_json(out_dir / "all_metrics.json", all_summary)
     print(json.dumps(all_summary, indent=2)[:4000])