Spaces:

Adive01
/

SummaryGenerator

Running

App Files Files Community

Adive01 commited on 14 days ago

Commit

fecd2b2

verified ·

1 Parent(s): 6aef09e

Upload mlplo/compare.py with huggingface_hub

Browse files

Files changed (1) hide show

mlplo/compare.py +222 -0

mlplo/compare.py ADDED Viewed

	@@ -0,0 +1,222 @@

+from __future__ import annotations
+import argparse
+import logging
+from pathlib import Path
+import evaluate
+import numpy as np
+import torch
+from datasets import load_from_disk
+from transformers import AutoModelForSeq2SeqLM
+from .common import (
+    ARTIFACT_DIR,
+    DEFAULT_SUMMARY_COLUMN,
+    DEFAULT_TEXT_COLUMN,
+    ensure_project_dirs,
+    load_tokenizer,
+    maybe_limit_split,
+    resolve_model_reference,
+    validate_model_dir,
+)
+LOGGER = logging.getLogger(__name__)
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(
+        description="Compare two models side-by-side on a test set."
+    )
+    parser.add_argument("--model-a", required=True, help="Path to Model A checkpoint.")
+    parser.add_argument("--model-b", required=True, help="Path to Model B checkpoint.")
+    parser.add_argument(
+        "--dataset-dir", required=True, help="Prepared dataset directory."
+    )
+    parser.add_argument("--split", default="test")
+    parser.add_argument("--max-samples", type=int, default=20)
+    parser.add_argument("--text-column", default=DEFAULT_TEXT_COLUMN)
+    parser.add_argument("--summary-column", default=DEFAULT_SUMMARY_COLUMN)
+    parser.add_argument(
+        "--output-file", default=str(ARTIFACT_DIR / "comparison.html")
+    )
+    return parser.parse_args()
+@torch.inference_mode()
+def generate_summaries(
+    model_path: str, dataset, text_col: str, device: torch.device
+) -> list[str]:
+    ref = resolve_model_reference(model_path)
+    validate_model_dir(ref)
+    LOGGER.info(f"Loading {ref}...")
+    tokenizer = load_tokenizer(ref)
+    model = AutoModelForSeq2SeqLM.from_pretrained(ref).to(device)
+    model.eval()
+    predictions = []
+    for item in dataset:
+        text = item[text_col]
+        inputs = tokenizer(
+            text, return_tensors="pt", truncation=True, max_length=512
+        ).to(device)
+        out = model.generate(**inputs, max_length=128, num_beams=4)
+        pred = tokenizer.decode(out[0], skip_special_tokens=True).strip()
+        predictions.append(pred)
+    del model
+    torch.cuda.empty_cache()
+    return predictions
+def score_predictions(predictions: list[str], references: list[str]) -> dict:
+    rouge = evaluate.load("rouge")
+    r_res = rouge.compute(
+        predictions=predictions, references=references, use_stemmer=True
+    )
+    from bert_score import score as bert_score_fn
+    safe_preds = [p if p.strip() else "..." for p in predictions]
+    safe_refs = [r if r.strip() else "..." for r in references]
+    LOGGER.info("Computing BERTScore...")
+    _, _, f1 = bert_score_fn(safe_preds, safe_refs, lang="en", verbose=False)
+    return {
+        "rouge1": r_res["rouge1"],
+        "rouge2": r_res["rouge2"],
+        "rougeL": r_res["rougeL"],
+        "bertscore": float(f1.mean().item()),
+    }
+def generate_html(
+    model_a_name: str,
+    model_b_name: str,
+    scores_a: dict,
+    scores_b: dict,
+    dataset,
+    preds_a: list[str],
+    preds_b: list[str],
+    text_col: str,
+    sum_col: str,
+) -> str:
+    html = f"""
+    <!DOCTYPE html>
+    <html>
+    <head>
+        <title>Model Comparison</title>
+        <style>
+            body {{ font-family: sans-serif; margin: 40px; color: #333; }}
+            table {{ border-collapse: collapse; width: 100%; margin-bottom: 30px; }}
+            th, td {{ border: 1px solid #ddd; padding: 12px; text-align: left; vertical-align: top; }}
+            th {{ background-color: #f8f9fa; font-weight: bold; }}
+            .better {{ background-color: #e8f5e9; font-weight: bold; color: #2e7d32; }}
+            .source-col {{ width: 30%; font-size: 0.9em; color: #555; }}
+            .ref-col {{ width: 20%; font-size: 0.9em; background: #fafafa; }}
+            .pred-col {{ width: 25%; }}
+        </style>
+    </head>
+    <body>
+        <h1>Model Comparison</h1>
+        <h2>Aggregate Scores</h2>
+        <table>
+            <tr>
+                <th>Metric</th>
+                <th>Model A: {model_a_name}</th>
+                <th>Model B: {model_b_name}</th>
+            </tr>
+    """
+    for k in ["rouge1", "rouge2", "rougeL", "bertscore"]:
+        va = scores_a[k]
+        vb = scores_b[k]
+        ca = "better" if va >= vb else ""
+        cb = "better" if vb > va else ""
+        html += f"""
+            <tr>
+                <td><strong>{k.upper()}</strong></td>
+                <td class="{ca}">{va:.4f}</td>
+                <td class="{cb}">{vb:.4f}</td>
+            </tr>
+        """
+    html += """
+        </table>
+        <h2>Side-by-Side Predictions</h2>
+        <table>
+            <tr>
+                <th>Source</th>
+                <th>Reference</th>
+                <th>Model A</th>
+                <th>Model B</th>
+            </tr>
+    """
+    for i, item in enumerate(dataset):
+        html += f"""
+            <tr>
+                <td class="source-col">{item[text_col]}</td>
+                <td class="ref-col">{item[sum_col]}</td>
+                <td class="pred-col">{preds_a[i]}</td>
+                <td class="pred-col">{preds_b[i]}</td>
+            </tr>
+        """
+    html += """
+        </table>
+    </body>
+    </html>
+    """
+    return html
+def main() -> None:
+    logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
+    args = parse_args()
+    ensure_project_dirs()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    LOGGER.info(f"Loading dataset {args.dataset_dir} (split: {args.split})...")
+    dataset = load_from_disk(args.dataset_dir)[args.split]
+    dataset = maybe_limit_split(dataset, args.max_samples)
+    refs = [item[args.summary_column] for item in dataset]
+    LOGGER.info("--- Processing Model A ---")
+    preds_a = generate_summaries(args.model_a, dataset, args.text_column, device)
+    scores_a = score_predictions(preds_a, refs)
+    LOGGER.info("--- Processing Model B ---")
+    preds_b = generate_summaries(args.model_b, dataset, args.text_column, device)
+    scores_b = score_predictions(preds_b, refs)
+    name_a = Path(args.model_a).name
+    name_b = Path(args.model_b).name
+    LOGGER.info("Generating HTML report...")
+    html = generate_html(
+        name_a,
+        name_b,
+        scores_a,
+        scores_b,
+        dataset,
+        preds_a,
+        preds_b,
+        args.text_column,
+        args.summary_column,
+    )
+    out_file = Path(args.output_file)
+    out_file.parent.mkdir(parents=True, exist_ok=True)
+    out_file.write_text(html, encoding="utf-8")
+    LOGGER.info(f"Comparison report written to {out_file.absolute()}")
+if __name__ == "__main__":
+    main()