# /// script
# dependencies = [
#   "trl>=0.12.0",
#   "peft>=0.13.0",
#   "transformers>=4.45.0",
#   "datasets>=3.0.0",
#   "accelerate>=1.0.0",
#   "trackio",
# ]
# ///
"""Lexwell SFT — Qwen2.5-3B-Instruct fine-tune on IRAC contract-review corpus.

Memory-tuned for a10g-large (24 GB): gradient checkpointing + batch 2 + grad-accum 2
keeps effective batch at 4 and total steps at 200 across 10 epochs.
"""
from datasets import load_dataset
from peft import LoraConfig
from trl import SFTTrainer, SFTConfig

BASE  = "Qwen/Qwen2.5-3B-Instruct"
DS    = "Curious-PM/lexwell-contract-irac"
OUT   = "Curious-PM/lexwell-contract-irac-qwen2.5-3b-lora"

print(f"Loading dataset {DS}...")
ds = load_dataset(DS, data_files="lexwell_v2.jsonl", split="train")
print(f"Loaded {len(ds)} rows")

trainer = SFTTrainer(
    model=BASE,
    train_dataset=ds,
    peft_config=LoraConfig(
        r=16,
        lora_alpha=32,
        task_type="CAUSAL_LM",
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    ),
    args=SFTConfig(
        output_dir="lexwell-lora",
        num_train_epochs=10,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=2,
        learning_rate=2e-4,
        max_length=2048,
        logging_steps=10,
        eval_strategy="no",
        save_strategy="no",
        bf16=True,
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        optim="adamw_torch_fused",
        dataloader_pin_memory=False,
        push_to_hub=True,
        hub_model_id=OUT,
        hub_strategy="end",
        report_to="none",
    ),
)

print("Starting training...")
trainer.train()
print("Training complete. Pushing to Hub...")
trainer.push_to_hub()
print(f"\nAdapter pushed: https://huggingface.co/{OUT}")