train_final.py

"""
SciHigh-2026 Subtask 1 (Research Highlight Generation) -- final training recipe.

Fine-tunes facebook/bart-large-cnn on the expanded MixSub-SciHigh training pool
to generate short research highlights from a paper's abstract. This is the
exact recipe that produced the submitted result:

    ROUGE-1=36.22%  ROUGE-2=13.17%  ROUGE-L=24.37%  ROUGE-Lsum=24.36%
    METEOR=32.48%   BERTScore-F1=87.49%
    (vs. the FIRE-2025 winning baseline of 23.45% ROUGE-L)

This script is a clean, standalone extraction of the winning configuration --
it deliberately does NOT include the large space of experiments (backbone
sweeps, input-feature engineering, decode-time reranking/logit-bias/checkpoint
averaging, proxy-scale fast iteration mode, etc.) that were tried and used to
arrive at this recipe. Those all live in the project's internal experiment
history; none of them are part of what actually produced the result above.

Usage:
    pip install torch transformers accelerate datasets evaluate rouge_score \
        sentencepiece bert-score nltk
    python train_final.py --data_dir /path/to/data --output_dir ./output

Expects --data_dir to contain:
    train_expanded.csv  (Filename, Abstract, Highlights -- 15,960 rows: the
                          official 10,000-row MixSub-SciHigh train split plus
                          5,960 additional real pairs recovered from the
                          dataset's original source release, leakage-checked
                          against val/test by exact Abstract-text match)
    val.csv             (Filename, Abstract, Highlights -- 1,985 rows, the
                          official held-out validation split, used only for
                          per-epoch monitoring/checkpoint selection here)
    test.csv             (Filename, Abstract -- 1,840 rows, official masked
                          test split, for the submission predictions)
"""
import argparse
import json
import os

import evaluate
import nltk
import numpy as np
import pandas as pd
import torch
from datasets import Dataset
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    EarlyStoppingCallback,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
)

MODEL_NAME = "facebook/bart-large-cnn"
MAX_INPUT_LEN = 512
MAX_TARGET_LEN = 100  # matches the FIRE-2025 baseline recipe's output budget
BATCH_SIZE = 2
LEARNING_RATE = 2e-5
NUM_BEAMS = 4

# Epoch ceiling, not a fixed schedule: bart-large-cnn converges fast on this
# task (a small-scale proxy run reached near-final quality in ~600-1,200
# gradient steps, a fraction of one full epoch's ~7,980 steps at batch_size=2
# over the 15,960-row pool). load_best_model_at_end + EarlyStoppingCallback
# below let the run self-terminate rather than committing to a fixed count.
# In the actual run that produced the reported numbers, training stopped
# after epoch 3 (2 consecutive non-improving epochs), and the checkpoint from
# epoch 1 -- the true best on val ROUGE-L -- was the one restored and
# evaluated/submitted.
EPOCH_CEILING = 4
EARLY_STOPPING_PATIENCE = 2

for _pkg in ["wordnet", "punkt_tab", "omw-1.4"]:
    try:
        nltk.download(_pkg, quiet=True)
    except Exception as e:  # noqa: BLE001
        print(f"warning: failed to download nltk resource '{_pkg}': {e}")


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--data_dir", required=True)
    p.add_argument("--output_dir", default="output")
    p.add_argument("--epochs", type=int, default=EPOCH_CEILING)
    p.add_argument("--skip_bertscore", action="store_true", help="BERTScore eval downloads its own scoring model; skip for a fast local check")
    return p.parse_args()


def to_hf_dataset(df, has_target):
    d = {"Abstract": df["Abstract"].tolist()}
    if has_target:
        d["Highlights"] = df["Highlights"].tolist()
    return Dataset.from_dict(d)


def make_preprocess_fn(tokenizer):
    def preprocess(batch):
        model_inputs = tokenizer(batch["Abstract"], max_length=MAX_INPUT_LEN, truncation=True)
        labels = tokenizer(text_target=batch["Highlights"], max_length=MAX_TARGET_LEN, truncation=True)
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs

    return preprocess


def main():
    args = parse_args()
    os.makedirs(args.output_dir, exist_ok=True)
    device = "cuda" if torch.cuda.is_available() else "cpu"

    train_df = pd.read_csv(os.path.join(args.data_dir, "train_expanded.csv"))
    val_df = pd.read_csv(os.path.join(args.data_dir, "val.csv"))
    test_df = pd.read_csv(os.path.join(args.data_dir, "test.csv"))
    print(f"[train_final] device={device} train={len(train_df)} val={len(val_df)} test={len(test_df)}")

    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME).to(device)

    # Decode-time settings: bart-large-cnn already ships no_repeat_ngram_size=3
    # in its own generation_config, but it's set explicitly here rather than
    # left implicit, so the decision doesn't silently depend on that shipped
    # default surviving a future transformers/model-card change. Two other
    # settings that were tuned during small-scale proxy experiments
    # (repetition_penalty=1.5, min_new_tokens=15) were fixes for degenerate
    # repetition-collapse in a severely undertrained checkpoint -- this
    # full-scale, fully-converged model doesn't exhibit that failure mode, so
    # those are deliberately left untouched at bart-large-cnn's own defaults
    # (repetition_penalty=1.0, min_length=56) rather than carried over.
    model.generation_config.no_repeat_ngram_size = 3
    model.generation_config.max_length = MAX_TARGET_LEN

    train_ds = to_hf_dataset(train_df, has_target=True)
    val_ds = to_hf_dataset(val_df, has_target=True)
    preprocess = make_preprocess_fn(tokenizer)
    train_tok = train_ds.map(preprocess, batched=True, remove_columns=train_ds.column_names)
    val_tok = val_ds.map(preprocess, batched=True, remove_columns=val_ds.column_names)
    collator = DataCollatorForSeq2Seq(tokenizer, model=model)

    rouge = evaluate.load("rouge")
    meteor = evaluate.load("meteor")

    def compute_metrics(eval_preds):
        preds, labels = eval_preds
        if isinstance(preds, tuple):
            preds = preds[0]
        preds = np.where(preds != -100, preds, tokenizer.pad_token_id)
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)

        result = rouge.compute(predictions=decoded_preds, references=decoded_labels)
        result = {f"rouge_{k}": v for k, v in result.items()}
        result["meteor"] = meteor.compute(predictions=decoded_preds, references=decoded_labels)["meteor"]
        return result

    training_args = Seq2SeqTrainingArguments(
        output_dir=os.path.join(args.output_dir, "checkpoints"),
        num_train_epochs=args.epochs,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        learning_rate=LEARNING_RATE,
        label_smoothing_factor=0.0,
        warmup_ratio=0.0,
        # Adafactor's factored second-moment estimates avoid the full-size
        # exp_avg/exp_avg_sq buffers that OOM'd a 16GB T4 with plain Adam; it's
        # also what the original PEGASUS/BART pretraining used.
        optim="adafactor",
        predict_with_generate=True,
        generation_max_length=MAX_TARGET_LEN,
        generation_num_beams=NUM_BEAMS,
        eval_strategy="epoch",
        save_strategy="epoch",
        save_total_limit=1,
        load_best_model_at_end=True,
        metric_for_best_model="rouge_rougeL",
        fp16=(device == "cuda"),
        logging_steps=50,
        report_to=[],
    )

    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_tok,
        eval_dataset=val_tok,
        data_collator=collator,
        compute_metrics=compute_metrics,
        callbacks=[EarlyStoppingCallback(early_stopping_patience=EARLY_STOPPING_PATIENCE)],
    )

    trainer.train()

    final_model_dir = os.path.join(args.output_dir, "model")
    trainer.save_model(final_model_dir)
    tokenizer.save_pretrained(final_model_dir)

    val_metrics = trainer.evaluate()
    print("[train_final] validation metrics:", val_metrics)

    if not args.skip_bertscore:
        from bert_score import score as bertscore

        val_preds = trainer.predict(val_tok)
        preds = np.where(val_preds.predictions != -100, val_preds.predictions, tokenizer.pad_token_id)
        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
        _, _, f1 = bertscore(decoded_preds, val_df["Highlights"].tolist(), lang="en", verbose=False)
        val_metrics["bertscore_f1"] = float(f1.mean())
        print("[train_final] bertscore_f1:", val_metrics["bertscore_f1"])

    with open(os.path.join(args.output_dir, "val_metrics.json"), "w") as f:
        json.dump(val_metrics, f, indent=2)

    # Generate the submission predictions on the masked test set.
    model.eval()
    gen_device = next(model.parameters()).device
    predictions = []
    batch_size = max(BATCH_SIZE, 8)
    abstracts = test_df["Abstract"].tolist()
    for i in range(0, len(abstracts), batch_size):
        batch = abstracts[i : i + batch_size]
        inputs = tokenizer(batch, max_length=MAX_INPUT_LEN, truncation=True, padding=True, return_tensors="pt").to(gen_device)
        with torch.no_grad():
            generated = model.generate(**inputs, max_length=MAX_TARGET_LEN, num_beams=NUM_BEAMS)
        predictions.extend(tokenizer.batch_decode(generated, skip_special_tokens=True))

    submission = pd.DataFrame({
        "Filename": test_df["Filename"],
        "Abstract": test_df["Abstract"],
        "Predicted_Highlights": predictions,
    })
    submission_path = os.path.join(args.output_dir, "Yushkk99_Task1_run1.csv")
    submission.to_csv(submission_path, index=False)
    print(f"[train_final] wrote submission to {submission_path}")


if __name__ == "__main__":
    main()