!pip install "unsloth[cu130-torch291] @ git+https://github.com/unslothai/unsloth.git"
!pip uninstall torch torchvision
!pip install torch torchvision --index-url https://download.pytorch.org/whl/cu130


from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset
import torch

max_seq_length = 2048
dtype = None
load_in_4bit = True

model, tokenizer = FastLanguageModel.from_pretrained(
    model_name="unsloth/gpt-oss-20b",
    max_seq_length=max_seq_length,
    dtype=dtype,
    load_in_4bit=load_in_4bit,
)

model = FastLanguageModel.get_peft_model(
    model,
    r=16,
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    lora_alpha=32,
    lora_dropout=0.05,
    bias="none",
    use_gradient_checkpointing="unsloth",
    random_state=3407,
    use_rslora=True,
)

dataset = load_dataset("AlicanKiraz0/Agentic-Chain-of-Thought-Coding-SFT-Dataset", split="train")

def formatting_prompts_func(examples):
    output_texts = []
    for i in range(len(examples['user'])):
        text = (
            f"### System:\n{examples['system'][i]}\n\n"
            f"### User:\n{examples['user'][i]}\n\n"
            f"### Assistant:\n{examples['assistant'][i]}{tokenizer.eos_token}"
        )
        output_texts.append(text)
    return output_texts

trainer = SFTTrainer(
    model=model,
    tokenizer=tokenizer,
    train_dataset=dataset,
    formatting_func=formatting_prompts_func,
    max_seq_length=max_seq_length,
    packing=False,
    args=TrainingArguments(
        per_device_train_batch_size=1,
        gradient_accumulation_steps=8,
        num_train_epochs=3,
        learning_rate=2e-5,
        bf16=torch.cuda.is_bf16_supported(),
        fp16=not torch.cuda.is_bf16_supported(),
        optim="adamw_8bit",
        output_dir="gpt-oss-20b-finetuned",
        report_to="none",
    ),
)

trainer.train()