import os
#os.environ['HF_HUB_OFFLINE'] = '1'
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
os.environ["HF_DEACTIVATE_ASYNC_LOAD"]="1"
 

from unsloth import FastLanguageModel
import torch
from datasets import load_dataset
from trl import SFTTrainer, SFTConfig

max_seq_length = 2048

dataset_path = r"E:\Unsloth\datasets\train.jsonl"
dataset = load_dataset("json", data_files=dataset_path, split="train")
 


model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = "trohrbaugh/Qwen3.5-9B-heretic-v2",
    max_seq_length = max_seq_length,
    load_in_4bit = True,
    local_files_only = True,
    trust_remote_code = True,
)

def formatting_prompts_func(examples):
    convos = examples["messages"]
    texts = [tokenizer.apply_chat_template(convo, tokenize=False, add_generation_prompt=False) for convo in convos]
    return { "text" : texts }

dataset = dataset.map(formatting_prompts_func, batched=True)

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 = 16,
    lora_dropout = 0,
    bias = "none",
    use_gradient_checkpointing = "unsloth",
    random_state = 3407,
    max_seq_length = max_seq_length,
)
sft_config = SFTConfig(
    dataset_text_field="text",
    max_seq_length = max_seq_length,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=32,
    warmup_steps=2,
    num_train_epochs=3,
    learning_rate=2e-4,
    logging_steps=1,
    output_dir="outputs_qwen35",
    optim="adamw_8bit",
    seed=3407,
    dataset_num_proc=None,
)

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    args = sft_config,
)

trainer.train()