#!/usr/bin/env python3 from __future__ import annotations import argparse import hashlib import json import os from datetime import datetime, timezone from pathlib import Path import evaluate import numpy as np import torch from datasets import DatasetDict, load_dataset from sklearn.metrics import classification_report, confusion_matrix from transformers import ( AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, EarlyStoppingCallback, Trainer, TrainingArguments, ) DEFAULT_MODEL_NAME = "allenai/scibert_scivocab_uncased" DEFAULT_PROJECT_ROOT = Path(__file__).resolve().parents[1] DEFAULT_DATA_DIR = DEFAULT_PROJECT_ROOT / "data" / "processed" DEFAULT_OUTPUT_DIR = DEFAULT_PROJECT_ROOT / "artifacts" / "scibert_topics12" DEFAULT_MAX_LENGTH = 256 DEFAULT_TRAIN_BATCH_SIZE = 4 DEFAULT_EVAL_BATCH_SIZE = 8 DEFAULT_GRAD_ACCUM_STEPS = 8 DEFAULT_NUM_EPOCHS = 12 DEFAULT_LEARNING_RATE = 2e-5 DEFAULT_WEIGHT_DECAY = 0.01 DEFAULT_WARMUP_RATIO = 0.1 DEFAULT_LABEL_SMOOTHING = 0.05 DEFAULT_TITLE_ONLY_PROB = 0.2 DEFAULT_EARLY_STOPPING_PATIENCE = 3 DEFAULT_SEED = 42 AUTO_RESUME_TOKEN = "auto" MODEL_FEATURE_COLUMNS = {"input_ids", "attention_mask", "token_type_ids", "labels"} def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Train a local article topic classifier on the prepared CSV dataset." ) parser.add_argument( "--data-dir", type=Path, default=DEFAULT_DATA_DIR, help="Directory with train.csv / validation.csv / test.csv / dataset_meta.json.", ) parser.add_argument( "--output-dir", type=Path, default=DEFAULT_OUTPUT_DIR, help="Directory where checkpoints and final artifacts will be stored.", ) parser.add_argument( "--model-name", default=DEFAULT_MODEL_NAME, help="HF model name to fine-tune.", ) parser.add_argument( "--max-length", type=int, default=DEFAULT_MAX_LENGTH, help="Tokenizer max_length.", ) parser.add_argument( "--per-device-train-batch-size", type=int, default=DEFAULT_TRAIN_BATCH_SIZE, help="Micro-batch size on one GPU.", ) parser.add_argument( "--per-device-eval-batch-size", type=int, default=DEFAULT_EVAL_BATCH_SIZE, help="Eval micro-batch size on one GPU.", ) parser.add_argument( "--gradient-accumulation-steps", type=int, default=DEFAULT_GRAD_ACCUM_STEPS, help="Gradient accumulation steps.", ) parser.add_argument( "--num-train-epochs", type=float, default=DEFAULT_NUM_EPOCHS, help="Maximum number of epochs.", ) parser.add_argument( "--learning-rate", type=float, default=DEFAULT_LEARNING_RATE, help="AdamW learning rate.", ) parser.add_argument( "--weight-decay", type=float, default=DEFAULT_WEIGHT_DECAY, help="AdamW weight decay.", ) parser.add_argument( "--warmup-ratio", type=float, default=DEFAULT_WARMUP_RATIO, help="Warmup ratio for cosine scheduler.", ) parser.add_argument( "--label-smoothing-factor", type=float, default=DEFAULT_LABEL_SMOOTHING, help="Label smoothing factor.", ) parser.add_argument( "--title-only-prob", type=float, default=DEFAULT_TITLE_ONLY_PROB, help="Fraction of train examples that will be converted to title-only inputs.", ) parser.add_argument( "--early-stopping-patience", type=int, default=DEFAULT_EARLY_STOPPING_PATIENCE, help="Number of eval rounds without improvement before stopping.", ) parser.add_argument( "--seed", type=int, default=DEFAULT_SEED, help="Random seed.", ) parser.add_argument( "--dataloader-num-workers", type=int, default=2, help="Number of dataloader workers.", ) parser.add_argument( "--expected-taxonomy-profile", default="topics12", help="Abort if prepared data was built with another taxonomy profile.", ) parser.add_argument( "--resume-from-checkpoint", default=None, help=( "Checkpoint path to resume from, or 'auto' to reuse the latest " "checkpoint already present in output-dir." ), ) parser.add_argument( "--skip-train", action="store_true", help="Skip fine-tuning and only run evaluation with the loaded model.", ) return parser.parse_args() def build_input_text(title: str, abstract: str, title_only: bool) -> str: title = (title or "").strip() abstract = (abstract or "").strip() if title_only or not abstract: return f"Title: {title}" return f"Title: {title}\n\nAbstract: {abstract}" def should_use_title_only(paper_id: str, probability: float, seed: int) -> bool: if probability <= 0: return False if probability >= 1: return True digest = hashlib.md5(f"{paper_id}:{seed}".encode("utf-8")).hexdigest() bucket = int(digest[:8], 16) / 0xFFFFFFFF return bucket < probability def load_metadata(path: Path) -> dict: if not path.exists(): raise FileNotFoundError(f"Missing dataset metadata: {path}") return json.loads(path.read_text(encoding="utf-8")) def resolve_resume_checkpoint(output_dir: Path, checkpoint_arg: str | None) -> str | None: if not checkpoint_arg: return None if checkpoint_arg != AUTO_RESUME_TOKEN: checkpoint_path = Path(checkpoint_arg) if not checkpoint_path.exists(): raise FileNotFoundError(f"Checkpoint does not exist: {checkpoint_path}") return str(checkpoint_path) checkpoints = sorted( output_dir.glob("checkpoint-*"), key=lambda path: int(path.name.split("-")[-1]), ) if not checkpoints: return None return str(checkpoints[-1]) def build_datasets( data_dir: Path, tokenizer, max_length: int, title_only_prob: float, seed: int, ) -> tuple[DatasetDict, DatasetDict]: data_files = { "train": str(data_dir / "train.csv"), "validation": str(data_dir / "validation.csv"), "test": str(data_dir / "test.csv"), } raw = load_dataset("csv", data_files=data_files) def add_text(example, split_name: str): title_only = split_name == "train" and should_use_title_only( example["paper_id"], title_only_prob, seed ) return { "model_text": build_input_text( example.get("title", ""), example.get("abstract", ""), title_only=title_only, ), "labels": int(example["label_id"]), } text_ready = DatasetDict() for split_name, split_dataset in raw.items(): text_ready[split_name] = split_dataset.map( lambda row, split_name=split_name: add_text(row, split_name), desc=f"Building model text for {split_name}", ) def tokenize_batch(batch): return tokenizer( batch["model_text"], truncation=True, max_length=max_length, ) tokenized = text_ready.map( tokenize_batch, batched=True, desc="Tokenizing", ) for split_name, split_dataset in tokenized.items(): removable = [ column_name for column_name in split_dataset.column_names if column_name not in MODEL_FEATURE_COLUMNS ] tokenized[split_name] = split_dataset.remove_columns(removable) return text_ready, tokenized def main() -> None: args = parse_args() os.environ["TOKENIZERS_PARALLELISM"] = "false" if torch.cuda.is_available(): torch.backends.cuda.matmul.allow_tf32 = True torch.backends.cudnn.allow_tf32 = True torch.set_float32_matmul_precision("high") metadata = load_metadata(args.data_dir / "dataset_meta.json") taxonomy_profile = metadata.get("taxonomy_profile") if args.expected_taxonomy_profile and taxonomy_profile != args.expected_taxonomy_profile: raise ValueError( f"Prepared dataset taxonomy_profile={taxonomy_profile!r}, " f"expected {args.expected_taxonomy_profile!r}" ) label_order = metadata["label_order"] label_display_names = metadata["label_display_names"] id2label = { idx: label_display_names[label] for idx, label in enumerate(label_order) } label2id = {display_name: idx for idx, display_name in id2label.items()} output_dir = args.output_dir output_dir.mkdir(parents=True, exist_ok=True) resume_checkpoint = resolve_resume_checkpoint( output_dir=output_dir, checkpoint_arg=args.resume_from_checkpoint, ) tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_fast=True) text_ready, tokenized = build_datasets( data_dir=args.data_dir, tokenizer=tokenizer, max_length=args.max_length, title_only_prob=args.title_only_prob, seed=args.seed, ) model = AutoModelForSequenceClassification.from_pretrained( args.model_name, num_labels=len(label_order), id2label=id2label, label2id=label2id, ) accuracy_metric = evaluate.load("accuracy") f1_metric = evaluate.load("f1") def compute_metrics(eval_pred): predictions, labels = eval_pred if predictions.ndim > 1: predictions = predictions.argmax(axis=-1) return { "accuracy": accuracy_metric.compute( predictions=predictions, references=labels )["accuracy"], "macro_f1": f1_metric.compute( predictions=predictions, references=labels, average="macro" )["f1"], } def preprocess_logits_for_metrics(logits, labels): if isinstance(logits, tuple): logits = logits[0] return logits.argmax(dim=-1) use_bf16 = torch.cuda.is_available() and torch.cuda.is_bf16_supported() use_fp16 = torch.cuda.is_available() and not use_bf16 training_args = TrainingArguments( output_dir=str(output_dir), do_train=True, do_eval=True, eval_strategy="epoch", save_strategy="epoch", logging_strategy="steps", logging_steps=50, save_total_limit=2, per_device_train_batch_size=args.per_device_train_batch_size, per_device_eval_batch_size=args.per_device_eval_batch_size, gradient_accumulation_steps=args.gradient_accumulation_steps, num_train_epochs=args.num_train_epochs, learning_rate=args.learning_rate, lr_scheduler_type="cosine", warmup_ratio=args.warmup_ratio, weight_decay=args.weight_decay, label_smoothing_factor=args.label_smoothing_factor, bf16=use_bf16, fp16=use_fp16, tf32=True if torch.cuda.is_available() else None, gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False}, eval_accumulation_steps=8, torch_empty_cache_steps=100, max_grad_norm=1.0, metric_for_best_model="macro_f1", greater_is_better=True, load_best_model_at_end=True, seed=args.seed, data_seed=args.seed, dataloader_num_workers=args.dataloader_num_workers, report_to="none", save_only_model=False, remove_unused_columns=True, disable_tqdm=False, logging_first_step=True, ) data_collator = DataCollatorWithPadding(tokenizer=tokenizer, pad_to_multiple_of=8) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized["train"], eval_dataset=tokenized["validation"], data_collator=data_collator, processing_class=tokenizer, compute_metrics=compute_metrics, preprocess_logits_for_metrics=preprocess_logits_for_metrics, callbacks=[ EarlyStoppingCallback( early_stopping_patience=args.early_stopping_patience ) ], ) run_config = { "prepared_at_utc": datetime.now(timezone.utc).isoformat(), "model_name": args.model_name, "max_length": args.max_length, "per_device_train_batch_size": args.per_device_train_batch_size, "per_device_eval_batch_size": args.per_device_eval_batch_size, "gradient_accumulation_steps": args.gradient_accumulation_steps, "num_train_epochs": args.num_train_epochs, "learning_rate": args.learning_rate, "weight_decay": args.weight_decay, "warmup_ratio": args.warmup_ratio, "label_smoothing_factor": args.label_smoothing_factor, "title_only_prob": args.title_only_prob, "early_stopping_patience": args.early_stopping_patience, "seed": args.seed, "resume_from_checkpoint": resume_checkpoint, "use_bf16": use_bf16, "use_fp16": use_fp16, "taxonomy_profile": taxonomy_profile, "num_labels": len(label_order), "label_order": label_order, } run_config_path = output_dir / ( "eval_run_config.json" if args.skip_train else "run_config.json" ) run_config_path.write_text( json.dumps(run_config, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) if args.skip_train: train_result_metrics = {"skipped_train": True} else: train_result = trainer.train(resume_from_checkpoint=resume_checkpoint) train_result_metrics = train_result.metrics trainer.save_model() tokenizer.save_pretrained(output_dir) metrics = { "train": trainer.evaluate(tokenized["train"], metric_key_prefix="train"), "validation": trainer.evaluate( tokenized["validation"], metric_key_prefix="validation" ), "test_full": trainer.evaluate(tokenized["test"], metric_key_prefix="test_full"), } title_only_test = text_ready["test"].map( lambda row: { "model_text": build_input_text( row.get("title", ""), row.get("abstract", ""), title_only=True, ) }, desc="Building title-only test view", ) tokenized_title_only_test = title_only_test.map( lambda batch: tokenizer( batch["model_text"], truncation=True, max_length=args.max_length, ), batched=True, desc="Tokenizing title-only test view", ) removable = [ column_name for column_name in tokenized_title_only_test.column_names if column_name not in MODEL_FEATURE_COLUMNS ] tokenized_title_only_test = tokenized_title_only_test.remove_columns(removable) metrics["test_title_only"] = trainer.evaluate( tokenized_title_only_test, metric_key_prefix="test_title_only", ) predictions = trainer.predict(tokenized["test"], metric_key_prefix="predict_test_full") full_pred_ids = predictions.predictions if full_pred_ids.ndim > 1: full_pred_ids = full_pred_ids.argmax(axis=-1) full_report = classification_report( predictions.label_ids, full_pred_ids, target_names=[id2label[idx] for idx in range(len(label_order))], digits=4, output_dict=True, zero_division=0, ) full_confusion = confusion_matrix( predictions.label_ids, full_pred_ids, ).tolist() title_only_predictions = trainer.predict( tokenized_title_only_test, metric_key_prefix="predict_test_title_only", ) title_only_pred_ids = title_only_predictions.predictions if title_only_pred_ids.ndim > 1: title_only_pred_ids = title_only_pred_ids.argmax(axis=-1) title_only_report = classification_report( title_only_predictions.label_ids, title_only_pred_ids, target_names=[id2label[idx] for idx in range(len(label_order))], digits=4, output_dict=True, zero_division=0, ) title_only_confusion = confusion_matrix( title_only_predictions.label_ids, title_only_pred_ids, ).tolist() summary = { "train_runtime": train_result_metrics, "eval_metrics": metrics, "test_full_classification_report": full_report, "test_full_confusion_matrix": full_confusion, "test_title_only_classification_report": title_only_report, "test_title_only_confusion_matrix": title_only_confusion, } (output_dir / "metrics_summary.json").write_text( json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) (output_dir / "metrics_summary.txt").write_text( json.dumps(summary, ensure_ascii=False, indent=2) + "\n", encoding="utf-8", ) print(json.dumps(summary, ensure_ascii=False, indent=2)) if __name__ == "__main__": main()