""" 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 title and abstract. Inputs are ' | <abstract>' with truncated abstracts repaired via DOI-verified Semantic Scholar recovery (see final/ pipeline). This is the exact recipe that produced the submitted result: ROUGE-1=39.30% ROUGE-2=15.32% ROUGE-L=26.37% METEOR=36.17% BERTScore-F1=88.05% (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_recovered_titled.csv (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_recovered_titled.csv (Filename, Abstract, Highlights -- 1,985 rows, the official held-out validation split, used only for per-epoch monitoring/checkpoint selection here) test_recovered_titled.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_recovered_titled.csv")) val_df = pd.read_csv(os.path.join(args.data_dir, "val_recovered_titled.csv")) test_df = pd.read_csv(os.path.join(args.data_dir, "test_recovered_titled.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()