#!/usr/bin/env python3
"""LoRA fine-tune t5-smaller from JSONL rows containing input and target."""

from __future__ import annotations

import argparse

from datasets import load_dataset
from peft import LoraConfig, TaskType, get_peft_model, prepare_model_for_kbit_training
from transformers import (
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    DataCollatorForSeq2Seq,
    Seq2SeqTrainer,
    Seq2SeqTrainingArguments,
)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--train-file", required=True)
    parser.add_argument("--model", default="ShinpacheShimura/t5-smaller")
    parser.add_argument("--subfolder", default="optimized-flan-t5-small")
    parser.add_argument("--output-dir", default="t5-smaller-lora")
    parser.add_argument("--epochs", type=float, default=3.0)
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--learning-rate", type=float, default=2e-4)
    args = parser.parse_args()

    common = {"subfolder": args.subfolder} if args.subfolder else {}
    tokenizer = AutoTokenizer.from_pretrained(args.model, **common)
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model, device_map="auto", **common)
    model = prepare_model_for_kbit_training(model)
    model = get_peft_model(
        model,
        LoraConfig(
            task_type=TaskType.SEQ_2_SEQ_LM,
            target_modules=["q", "v"],
            r=8,
            lora_alpha=16,
            lora_dropout=0.05,
        ),
    )

    dataset = load_dataset("json", data_files=args.train_file, split="train")
    missing = {"input", "target"} - set(dataset.column_names)
    if missing:
        raise SystemExit(f"Missing fields: {', '.join(sorted(missing))}")

    def tokenize(batch):
        encoded = tokenizer(batch["input"], truncation=True, max_length=256)
        encoded["labels"] = tokenizer(
            text_target=batch["target"], truncation=True, max_length=128
        )["input_ids"]
        return encoded

    tokenized = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names)
    train_args = Seq2SeqTrainingArguments(
        output_dir=args.output_dir,
        num_train_epochs=args.epochs,
        per_device_train_batch_size=args.batch_size,
        gradient_accumulation_steps=4,
        learning_rate=args.learning_rate,
        logging_steps=10,
        save_strategy="epoch",
        report_to="none",
        fp16=True,
    )
    trainer = Seq2SeqTrainer(
        model=model,
        args=train_args,
        train_dataset=tokenized,
        processing_class=tokenizer,
        data_collator=DataCollatorForSeq2Seq(tokenizer, model=model),
    )
    trainer.train()
    trainer.save_model(args.output_dir)
    tokenizer.save_pretrained(args.output_dir)


if __name__ == "__main__":
    main()