finetune.py

#!/usr/bin/env python3
"""Fine-tune rubai-corrector-base on JSONL correction pairs."""

from __future__ import annotations

import argparse
import json
import random
from pathlib import Path
from typing import Any

import torch
from torch.utils.data import Dataset
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
    get_linear_schedule_with_warmup,
    set_seed,
)
from transformers.optimization import Adafactor


INPUT_PREFIX = "correct: "


def resolve_jsonl_path(path: Path) -> Path:
    if path.is_dir():
        candidate = path / "data.jsonl"
        if candidate.exists():
            return candidate
        raise FileNotFoundError(f"Directory {path} does not contain data.jsonl")
    return path


def load_records(path: Path) -> list[dict[str, Any]]:
    data_path = resolve_jsonl_path(path)
    records: list[dict[str, Any]] = []
    with data_path.open("r", encoding="utf-8") as handle:
        for line_num, line in enumerate(handle, start=1):
            line = line.strip()
            if not line:
                continue
            record = json.loads(line)
            if not isinstance(record.get("input"), str) or not isinstance(record.get("output"), str):
                raise ValueError(
                    f"{data_path}:{line_num} must contain string fields 'input' and 'output'"
                )
            records.append(record)
    if not records:
        raise ValueError(f"No records loaded from {data_path}")
    return records


def split_records(
    records: list[dict[str, Any]],
    validation_split: float,
    seed: int,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    if validation_split <= 0:
        return records, []
    if not 0 < validation_split < 1:
        raise ValueError("--validation-split must be between 0 and 1")
    items = records[:]
    random.Random(seed).shuffle(items)
    eval_size = max(1, int(len(items) * validation_split))
    return items[eval_size:], items[:eval_size]


class CorrectionDataset(Dataset):
    def __init__(
        self,
        records: list[dict[str, Any]],
        tokenizer,
        max_source_length: int,
        max_target_length: int,
    ):
        self.records = records
        self.tokenizer = tokenizer
        self.max_source_length = max_source_length
        self.max_target_length = max_target_length

    def __len__(self) -> int:
        return len(self.records)

    def __getitem__(self, index: int) -> dict[str, Any]:
        record = self.records[index]
        model_inputs = self.tokenizer(
            INPUT_PREFIX + record["input"],
            truncation=True,
            max_length=self.max_source_length,
        )
        labels = self.tokenizer(
            record["output"],
            truncation=True,
            max_length=self.max_target_length,
        )
        model_inputs["labels"] = labels["input_ids"]
        return model_inputs


class AdafactorSeq2SeqTrainer(Seq2SeqTrainer):
    def create_optimizer(self):
        if self.optimizer is None:
            self.optimizer = Adafactor(
                self.model.parameters(),
                lr=self.args.learning_rate,
                scale_parameter=False,
                relative_step=False,
                warmup_init=False,
                weight_decay=self.args.weight_decay,
            )
        return self.optimizer

    def create_scheduler(self, num_training_steps: int, optimizer=None):
        if self.lr_scheduler is None:
            actual_optimizer = optimizer if optimizer is not None else self.optimizer
            self.lr_scheduler = get_linear_schedule_with_warmup(
                actual_optimizer,
                num_warmup_steps=self.args.get_warmup_steps(num_training_steps),
                num_training_steps=num_training_steps,
            )
        return self.lr_scheduler


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model-path", type=Path, default=Path(__file__).resolve().parent)
    parser.add_argument("--train-file", type=Path, required=True)
    parser.add_argument("--eval-file", type=Path, default=None)
    parser.add_argument("--validation-split", type=float, default=0.0)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--max-source-length", type=int, default=512)
    parser.add_argument("--max-target-length", type=int, default=512)
    parser.add_argument("--learning-rate", type=float, default=5e-5)
    parser.add_argument("--weight-decay", type=float, default=0.01)
    parser.add_argument("--warmup-ratio", type=float, default=0.1)
    parser.add_argument("--num-train-epochs", type=float, default=2.0)
    parser.add_argument("--per-device-train-batch-size", type=int, default=16)
    parser.add_argument("--per-device-eval-batch-size", type=int, default=16)
    parser.add_argument("--gradient-accumulation-steps", type=int, default=4)
    parser.add_argument("--save-steps", type=int, default=500)
    parser.add_argument("--eval-steps", type=int, default=500)
    parser.add_argument("--logging-steps", type=int, default=50)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--bf16", action="store_true")
    parser.add_argument("--fp16", action="store_true")
    parser.add_argument("--gradient-checkpointing", action="store_true", default=True)
    parser.add_argument("--no-gradient-checkpointing", action="store_true")
    parser.add_argument("--resume-from-checkpoint", type=str, default=None)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    set_seed(args.seed)

    gradient_checkpointing = args.gradient_checkpointing and not args.no_gradient_checkpointing

    tokenizer = AutoTokenizer.from_pretrained(args.model_path)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_path)
    if gradient_checkpointing:
        model.gradient_checkpointing_enable()

    train_records = load_records(args.train_file)
    if args.eval_file is not None:
        eval_records = load_records(args.eval_file)
    else:
        train_records, eval_records = split_records(train_records, args.validation_split, args.seed)

    train_dataset = CorrectionDataset(
        train_records,
        tokenizer,
        max_source_length=args.max_source_length,
        max_target_length=args.max_target_length,
    )
    eval_dataset = None
    if eval_records:
        eval_dataset = CorrectionDataset(
            eval_records,
            tokenizer,
            max_source_length=args.max_source_length,
            max_target_length=args.max_target_length,
        )

    data_collator = DataCollatorForSeq2Seq(
        tokenizer=tokenizer,
        model=model,
        label_pad_token_id=-100,
        pad_to_multiple_of=8 if torch.cuda.is_available() else None,
    )

    training_args = Seq2SeqTrainingArguments(
        output_dir=str(args.output_dir),
        learning_rate=args.learning_rate,
        weight_decay=args.weight_decay,
        warmup_ratio=args.warmup_ratio,
        num_train_epochs=args.num_train_epochs,
        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,
        logging_steps=args.logging_steps,
        save_steps=args.save_steps,
        eval_steps=args.eval_steps,
        evaluation_strategy="steps" if eval_dataset is not None else "no",
        save_strategy="steps",
        save_total_limit=2,
        predict_with_generate=False,
        report_to=[],
        bf16=args.bf16,
        fp16=args.fp16 and not args.bf16,
        gradient_checkpointing=gradient_checkpointing,
        dataloader_num_workers=2,
        remove_unused_columns=False,
        seed=args.seed,
    )

    trainer = AdafactorSeq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        tokenizer=tokenizer,
        data_collator=data_collator,
    )

    trainer.train(resume_from_checkpoint=args.resume_from_checkpoint)
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())