"""T5-small fine-tuning and inference for one-sentence tiny summaries. Training uses `short_ref -> tiny_ref` pairs. At inference, the pipeline feeds the generated short summary to the fine-tuned T5 checkpoint, which keeps tiny generation cheap and less exposed to the full legal source. """ from __future__ import annotations import argparse import inspect from dataclasses import dataclass from pathlib import Path import pandas as pd from src.utils import MODELS_DIR, RESULTS_DIR, get_logger, set_seed logger = get_logger(__name__) DEFAULT_MODEL_DIR = MODELS_DIR / "t5_tiny_summarizer" DEFAULT_BASE_MODEL = "t5-small" @dataclass(frozen=True) class TinyTrainConfig: input_path: Path = Path("data/multilexsum_clean.parquet") output_dir: Path = DEFAULT_MODEL_DIR source_column: str = "short_ref" target_column: str = "tiny_ref" base_model: str = DEFAULT_BASE_MODEL epochs: int = 3 learning_rate: float = 5e-5 batch_size: int = 4 max_source_length: int = 256 max_target_length: int = 48 max_train: int | None = None max_val: int | None = None seed: int = 42 class TinyT5Summarizer: """Inference wrapper for the fine-tuned tiny-summary model.""" def __init__( self, model_dir: str | Path = DEFAULT_MODEL_DIR, fallback_model: str = DEFAULT_BASE_MODEL, device: str | None = None, max_source_length: int = 256, max_new_tokens: int = 40, ) -> None: try: import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer except ImportError as exc: raise ImportError("Install transformers, torch, and sentencepiece before tiny inference.") from exc model_path = Path(model_dir) model_name = str(model_path) if (model_path / "config.json").exists() else fallback_model self.tokenizer = AutoTokenizer.from_pretrained(model_name) self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name) self.torch = torch self.max_source_length = max_source_length self.max_new_tokens = max_new_tokens if device is None: if torch.cuda.is_available(): device = "cuda" elif getattr(torch.backends, "mps", None) and torch.backends.mps.is_available(): device = "mps" else: device = "cpu" self.device = device self.model.to(device) self.model.eval() logger.info("Loaded tiny summarizer %s on %s", model_name, device) def summarize(self, short_summary: str) -> str: prompt = f"summarize: {short_summary.strip()}" inputs = self.tokenizer( prompt, return_tensors="pt", truncation=True, max_length=self.max_source_length, ) inputs = {k: v.to(self.device) for k, v in inputs.items()} with self.torch.no_grad(): output_ids = self.model.generate( **inputs, max_new_tokens=self.max_new_tokens, num_beams=4, no_repeat_ngram_size=2, length_penalty=0.8, ) return self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip() def _prepare_dataset(df: pd.DataFrame, source_column: str, target_column: str, max_rows: int | None): from datasets import Dataset work = df[[source_column, target_column]].dropna().rename( columns={source_column: "source", target_column: "target"} ) if max_rows is not None: work = work.head(max_rows) return Dataset.from_pandas(work.reset_index(drop=True)) def _training_args_kwargs(config: TinyTrainConfig) -> dict: from transformers import Seq2SeqTrainingArguments params = inspect.signature(Seq2SeqTrainingArguments.__init__).parameters kwargs = { "output_dir": str(config.output_dir), "learning_rate": config.learning_rate, "per_device_train_batch_size": config.batch_size, "per_device_eval_batch_size": config.batch_size, "num_train_epochs": config.epochs, "weight_decay": 0.01, "predict_with_generate": True, "logging_steps": 25, "save_strategy": "epoch", "save_total_limit": 2, "report_to": [], "seed": config.seed, } if "eval_strategy" in params: kwargs["eval_strategy"] = "epoch" else: kwargs["evaluation_strategy"] = "epoch" return kwargs def train_tiny_model(config: TinyTrainConfig) -> pd.DataFrame: """Fine-tune T5-small and write the checkpoint plus validation-loss logs.""" try: from transformers import ( AutoModelForSeq2SeqLM, AutoTokenizer, DataCollatorForSeq2Seq, Seq2SeqTrainer, Seq2SeqTrainingArguments, ) except ImportError as exc: raise ImportError("Install transformers, torch, datasets, and sentencepiece before training.") from exc set_seed(config.seed) df = pd.read_parquet(config.input_path) train_df = df[df["split"] == "train"] val_df = df[df["split"] == "val"] train_ds = _prepare_dataset(train_df, config.source_column, config.target_column, config.max_train) val_ds = _prepare_dataset(val_df, config.source_column, config.target_column, config.max_val) tokenizer = AutoTokenizer.from_pretrained(config.base_model) model = AutoModelForSeq2SeqLM.from_pretrained(config.base_model) def preprocess(batch): inputs = [f"summarize: {x}" for x in batch["source"]] model_inputs = tokenizer( inputs, max_length=config.max_source_length, truncation=True, ) labels = tokenizer( text_target=batch["target"], max_length=config.max_target_length, truncation=True, ) model_inputs["labels"] = labels["input_ids"] return model_inputs tokenized_train = train_ds.map(preprocess, batched=True, remove_columns=train_ds.column_names) tokenized_val = val_ds.map(preprocess, batched=True, remove_columns=val_ds.column_names) args = Seq2SeqTrainingArguments(**_training_args_kwargs(config)) collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model) trainer_kwargs = { "model": model, "args": args, "train_dataset": tokenized_train, "eval_dataset": tokenized_val, "data_collator": collator, } trainer_params = inspect.signature(Seq2SeqTrainer.__init__).parameters if "tokenizer" in trainer_params: trainer_kwargs["tokenizer"] = tokenizer elif "processing_class" in trainer_params: trainer_kwargs["processing_class"] = tokenizer trainer = Seq2SeqTrainer(**trainer_kwargs) trainer.train() trainer.save_model(str(config.output_dir)) tokenizer.save_pretrained(str(config.output_dir)) logger.info("Saved fine-tuned T5 checkpoint -> %s", config.output_dir) history = pd.DataFrame(trainer.state.log_history) results_path = RESULTS_DIR / "t5_tiny_val_loss.csv" RESULTS_DIR.mkdir(parents=True, exist_ok=True) history.to_csv(results_path, index=False) logger.info("Wrote training log -> %s", results_path) if "eval_loss" in history.columns: import matplotlib.pyplot as plt curve = history.dropna(subset=["eval_loss"]) if not curve.empty: plt.figure(figsize=(6, 4)) plt.plot(curve["epoch"], curve["eval_loss"], marker="o") plt.xlabel("Epoch") plt.ylabel("Validation loss") plt.title("T5-small tiny-summary fine-tune") plt.tight_layout() plot_path = RESULTS_DIR / "t5_tiny_val_loss.png" plt.savefig(plot_path, dpi=160) plt.close() logger.info("Wrote validation-loss curve -> %s", plot_path) return history def build_arg_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description="Fine-tune or run T5-small tiny summarization.") parser.add_argument("--train", action="store_true", help="Fine-tune the model.") parser.add_argument("--input", default="data/multilexsum_clean.parquet") parser.add_argument("--output-dir", default=str(DEFAULT_MODEL_DIR)) parser.add_argument("--source-column", default="short_ref") parser.add_argument("--target-column", default="tiny_ref") parser.add_argument("--base-model", default=DEFAULT_BASE_MODEL) parser.add_argument("--epochs", type=int, default=3) parser.add_argument("--lr", type=float, default=5e-5) parser.add_argument("--batch-size", type=int, default=4) parser.add_argument("--max-train", type=int) parser.add_argument("--max-val", type=int) parser.add_argument("--text", help="Short summary to compress at inference time.") parser.add_argument("--device", choices=["cpu", "cuda", "mps"]) return parser def main() -> None: args = build_arg_parser().parse_args() if args.train: config = TinyTrainConfig( input_path=Path(args.input), output_dir=Path(args.output_dir), source_column=args.source_column, target_column=args.target_column, base_model=args.base_model, epochs=args.epochs, learning_rate=args.lr, batch_size=args.batch_size, max_train=args.max_train, max_val=args.max_val, ) train_tiny_model(config) return if not args.text: raise SystemExit("Provide --train or --text.") summarizer = TinyT5Summarizer(model_dir=args.output_dir, device=args.device) print(summarizer.summarize(args.text)) if __name__ == "__main__": main()