#!/usr/bin/env python3
"""Full fine-tune gemma-3-270m to reconstruct prompts from model outputs."""

from datasets import load_from_disk
from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForSeq2Seq

# --- Configuration ---
MODEL_NAME = "google/gemma-3-270m"
DATASET_PATH = "tokenized-dataset-plain"
OUTPUT_DIR = "checkpoints-270m-ft"
MAX_SEQ_LENGTH = 2048

# --- Training ---
BATCH_SIZE = 2
GRAD_ACCUM = 8
LEARNING_RATE = 5e-5
NUM_EPOCHS = 1
WARMUP_STEPS = 100
LOGGING_STEPS = 1
SAVE_STEPS = 100


def main():
    print("Loading pre-tokenized dataset...")
    dataset = load_from_disk(DATASET_PATH)
    print(f"Training examples: {len(dataset)}")

    print(f"Loading model: {MODEL_NAME}")
    model = AutoModelForCausalLM.from_pretrained(
        MODEL_NAME,
        torch_dtype="bfloat16",
        device_map="auto",
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    trainer = Trainer(
        model=model,
        processing_class=tokenizer,
        train_dataset=dataset,
        data_collator=DataCollatorForSeq2Seq(tokenizer, padding=True),
        args=TrainingArguments(
            output_dir=OUTPUT_DIR,
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=GRAD_ACCUM,
            learning_rate=LEARNING_RATE,
            num_train_epochs=NUM_EPOCHS,
            warmup_steps=WARMUP_STEPS,
            logging_steps=LOGGING_STEPS,
            save_steps=SAVE_STEPS,
            bf16=True,
            seed=42,
            report_to="wandb",
            logging_strategy="steps",
            gradient_checkpointing=True,
        ),
    )

    print("Training...")
    trainer.train()

    print("Saving model...")
    trainer.save_model(f"{OUTPUT_DIR}/final")
    tokenizer.save_pretrained(f"{OUTPUT_DIR}/final")
    print("Done.")


if __name__ == "__main__":
    main()