from unsloth import FastLanguageModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, Mxfp4Config
import re
import torch.nn as nn
import bitsandbytes as bnb

model_id = "unsloth/gpt-oss-120b"
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_id,
    dtype = torch.bfloat16, # None for auto detection
    max_seq_length = 32786, # Choose any for long context!
    load_in_4bit = True,  # 4 bit quantization to reduce memory
    full_finetuning = False, # [NEW!] We have full finetuning now!
)
print(model)


PATTERNS = [
    r"\.self_attn\.(q_proj|k_proj|v_proj|o_proj)$",
    r"\.mlp\.experts\.gate_up_projs\.\d+$",        # every expert's up proj
    r"\.mlp\.experts\.down_projs\.\d+$",           # every expert's down proj
]
PATTERNS = [re.compile(p) for p in PATTERNS]

def is_target(name, module):
    if not isinstance(module, (nn.Linear, bnb.nn.Linear4bit)):
        return False
    return any(p.search(name) for p in PATTERNS)

target_modules = [name for name, mod in model.named_modules() if is_target(name, mod)]
print(f"{len(target_modules)} modules will get LoRA")

model = FastLanguageModel.get_peft_model(
    model,
    r = 32, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
    target_modules=target_modules,
    lora_alpha = 64,
    lora_dropout = 0, # Supports any, but = 0 is optimized
    bias = "none",    # Supports any, but = "none" is optimized
    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
    use_gradient_checkpointing = "unsloth", # True or "unsloth" for very long context
    random_state = 3407,
    use_rslora = True,  # We support rank stabilized LoRA
    loftq_config = None, # And LoftQ
)


from datasets import load_from_disk
train = load_from_disk("data/OpenMathReasoning-harmony-deduplicated")

def tok(batch):
    enc = tokenizer(
        [p + c for p, c in zip(batch["prompt"], batch["completion"])],
        padding=False,
        truncation=False,
    )

    labels = []
    for p, c, ids in zip(batch["prompt"], batch["completion"], enc["input_ids"]):
        prompt_ids = tokenizer(p, add_special_tokens=False)["input_ids"]
        cutoff = len(prompt_ids)  # everything up to here gets masked out
        labels.append([-100] * cutoff + ids[cutoff:])
    enc["labels"] = labels
    return enc

train_tok = train.map(
    tok,
    batched=True,
    batch_size=1000,
    writer_batch_size=10_000,
    num_proc=24,
    remove_columns=train.column_names,
)

from transformers import TrainingArguments
from trl import SFTConfig, SFTTrainer


training_args = TrainingArguments(
    learning_rate=5.0e-5,
    gradient_checkpointing=True,
    num_train_epochs=1,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    warmup_ratio=0.03,
    lr_scheduler_type="cosine_with_min_lr",
    lr_scheduler_kwargs={"min_lr_rate": 0.1},
    output_dir="gpt-oss-120b-sft",
    report_to="none",
    bf16=True,
    fp16=False,
    optim = "paged_adamw_8bit",

    save_strategy="steps",          # or "epoch"
    save_steps=200,                 # pick a cadence that fits your run length
    save_total_limit=2,             # keep disk usage in check
    logging_steps=20,
)


from transformers import Trainer
from transformers.trainer_utils import get_last_checkpoint
from torch.nn.utils.rnn import pad_sequence

def collate(batch):
    input_ids = [torch.tensor(x["input_ids"], dtype=torch.long) for x in batch]
    attention_mask = [torch.tensor(x["attention_mask"], dtype=torch.long) for x in batch]
    labels = [torch.tensor(x["labels"], dtype=torch.long) for x in batch]

    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
    attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0)
    labels = pad_sequence(labels, batch_first=True, padding_value=-100)

    return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels}

trainer = Trainer(
    model=model,
    tokenizer=tokenizer,
    args=training_args,
    train_dataset=train_tok,
    data_collator=collate,
)

last_ckpt = get_last_checkpoint(training_args.output_dir)
if last_ckpt:
    trainer.train(resume_from_checkpoint=last_ckpt)
else:
    trainer.train()

model.save_pretrained_merged(training_args.output_dir, tokenizer, save_method="mxfp4")