Upload mlplo/eval.py with huggingface_hub

mlplo/eval.py

@@ -0,0 +1,189 @@
+from __future__ import annotations
+
+import argparse
+import logging
+from pathlib import Path
+
+import numpy as np
+from datasets import load_from_disk
+from transformers import (
+    AutoModelForSeq2SeqLM,
+    DataCollatorForSeq2Seq,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+)
+
+from .common import (
+    ARTIFACT_DIR,
+    DEFAULT_SUMMARY_COLUMN,
+    DEFAULT_TARGET_MAX_LENGTH,
+    DEFAULT_TEXT_COLUMN,
+    build_compute_metrics,
+    ensure_project_dirs,
+    existing_default_checkpoint,
+    load_tokenizer,
+    maybe_limit_split,
+    resolve_mixed_precision,
+    resolve_model_reference,
+    validate_model_dir,
+    write_json,
+    write_jsonl,
+)
+
+LOGGER = logging.getLogger(__name__)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Evaluate a fine-tuned summarization checkpoint."
+    )
+    parser.add_argument(
+        "--dataset-dir", required=True, help="Path produced by mlplo.data_cleaning."
+    )
+    parser.add_argument("--model-path", default=existing_default_checkpoint())
+    parser.add_argument(
+        "--split", default="test", choices=["train", "validation", "test"]
+    )
+    parser.add_argument("--text-column", default=DEFAULT_TEXT_COLUMN)
+    parser.add_argument("--summary-column", default=DEFAULT_SUMMARY_COLUMN)
+    parser.add_argument("--per-device-eval-batch-size", type=int, default=2)
+    parser.add_argument(
+        "--generation-max-length", type=int, default=DEFAULT_TARGET_MAX_LENGTH
+    )
+    parser.add_argument("--generation-num-beams", type=int, default=4)
+    parser.add_argument("--max-samples", type=int, default=None)
+    parser.add_argument("--preview-rows", type=int, default=5)
+    parser.add_argument(
+        "--include-bertscore",
+        action="store_true",
+        help=(
+            "Compute BERTScore F1 in addition to ROUGE. "
+            "Downloads a ~400 MB model on first use."
+        ),
+    )
+    parser.add_argument(
+        "--output-file", default=str(ARTIFACT_DIR / "eval_metrics.json")
+    )
+    parser.add_argument(
+        "--predictions-file", default=str(ARTIFACT_DIR / "sample_predictions.jsonl")
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
+    )
+    args = parse_args()
+    ensure_project_dirs()
+
+    if not args.model_path:
+        raise ValueError(
+            "No model path provided and no default checkpoint exists yet. "
+            "Train a model first with mlplo.train."
+        )
+
+    # ── Validate dataset path ─────────────────────────────────────────────────
+    dataset_path = Path(args.dataset_dir)
+    if not dataset_path.exists():
+        raise FileNotFoundError(f"Prepared dataset not found: {dataset_path}")
+
+    # ── Validate model directory ──────────────────────────────────────────────
+    model_reference = resolve_model_reference(args.model_path)
+    validate_model_dir(model_reference)
+
+    LOGGER.info("Loading dataset from %s", dataset_path)
+    tokenized_dataset = load_from_disk(str(dataset_path))
+
+    if args.split not in tokenized_dataset:
+        available = list(tokenized_dataset.keys())
+        raise KeyError(
+            f"Split '{args.split}' not found in dataset. Available: {available}"
+        )
+
+    evaluation_split = maybe_limit_split(
+        tokenized_dataset[args.split], args.max_samples
+    )
+
+    # ── Load model ────────────────────────────────────────────────────────────
+    LOGGER.info("Loading model from %s", model_reference)
+    tokenizer = load_tokenizer(model_reference)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_reference)
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model)
+    precision = resolve_mixed_precision()
+
+    if args.include_bertscore:
+        LOGGER.info(
+            "BERTScore enabled. A ~400 MB model will be downloaded on first use."
+        )
+
+    compute_metrics = build_compute_metrics(
+        tokenizer, include_bertscore=args.include_bertscore
+    )
+
+    temp_output_dir = ARTIFACT_DIR / "eval_tmp"
+    evaluation_args = Seq2SeqTrainingArguments(
+        output_dir=str(temp_output_dir),
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        predict_with_generate=True,
+        generation_max_length=args.generation_max_length,
+        generation_num_beams=args.generation_num_beams,
+        fp16=precision["fp16"],
+        bf16=precision["bf16"],
+        report_to="none",
+    )
+
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=evaluation_args,
+        processing_class=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    LOGGER.info("Running evaluation on split '%s'…", args.split)
+    prediction_output = trainer.predict(evaluation_split, metric_key_prefix=args.split)
+    metrics = prediction_output.metrics
+    write_json(args.output_file, metrics)
+    LOGGER.info("Metrics: %s", metrics)
+
+    # ── Decode predictions and write sample file ──────────────────────────────
+    generated_token_ids = prediction_output.predictions
+    if isinstance(generated_token_ids, tuple):
+        generated_token_ids = generated_token_ids[0]
+
+    generated_token_ids = np.asarray(generated_token_ids)
+    generated_token_ids = np.where(
+        generated_token_ids < 0, tokenizer.pad_token_id, generated_token_ids
+    )
+    decoded_predictions = tokenizer.batch_decode(
+        generated_token_ids, skip_special_tokens=True
+    )
+
+    # Guard against preview_rows exceeding available samples
+    n_preview = min(args.preview_rows, len(decoded_predictions), len(evaluation_split))
+    preview_rows = []
+    for index in range(n_preview):
+        row = evaluation_split[index]
+        prediction = decoded_predictions[index].strip()
+        record: dict = {
+            "source": row.get(args.text_column, ""),
+            "reference": row.get(args.summary_column, ""),
+            "prediction": prediction,
+        }
+        if not prediction:
+            record["empty_prediction"] = True
+            LOGGER.warning("Empty prediction at index %d.", index)
+        preview_rows.append(record)
+
+    write_jsonl(args.predictions_file, preview_rows)
+    LOGGER.info(
+        "Evaluation complete. Metrics → %s | Predictions → %s",
+        args.output_file,
+        args.predictions_file,
+    )
+
+
+if __name__ == "__main__":
+    main()