qwen lora

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import torch
from datasets import load_dataset
from transformers import Trainer, TrainingArguments, DataCollatorForCompletionOnlyLM
from unsloth import FastLanguageModel, is_bfloat16_supported

MODEL_NAME = "Qwen/Qwen3.5-4B-Base"
DATASET_PATH = os.path.join(os.path.dirname(__file__), '..', 'datasets', 'MyDataset.jsonl')
OUTPUT_DIR = os.path.join(os.path.dirname(__file__), '..', 'qwen_lora')
MAX_SEQ_LENGTH = 4096
LORA_R = 16
LORA_ALPHA = 32
LORA_DROPOUT = 0
BATCH_SIZE = 1
GRAD_ACCUM = 8
LEARNING_RATE = 2e-4
NUM_EPOCHS = 5
WARMUP_STEPS = 10

def main():
    model, tokenizer = FastLanguageModel.from_pretrained(
        MODEL_NAME,
        max_seq_length=MAX_SEQ_LENGTH,
        load_in_4bit=True,
        dtype=None,
        local_files_only=True,
    )

    model = FastLanguageModel.get_peft_model(
        model,
        r=LORA_R,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj",
                        "in_proj_qkv", "in_proj_z", "out_proj",
                        "in_proj_a", "in_proj_b"],
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        bias="none",
        use_gradient_checkpointing="unsloth",
    )
    model.print_trainable_parameters()

    text_tok = tokenizer.tokenizer
    text_tok.pad_token = text_tok.eos_token

    dataset = load_dataset("json", data_files=DATASET_PATH, split="train")

    def tokenize_fn(examples):
        result = text_tok(
            examples["text"],
            truncation=True,
            max_length=MAX_SEQ_LENGTH,
            padding=False,
        )
        return result

    dataset = dataset.map(tokenize_fn, batched=True, remove_columns=["text"])
    print(f"Examples: {len(dataset)}")

    response_template = "<|im_start|>assistant\n"

    data_collator = DataCollatorForCompletionOnlyLM(
        response_template=response_template,
        tokenizer=text_tok,
        mlm=False,
    )

    trainer = Trainer(
        model=model,
        args=TrainingArguments(
            output_dir=OUTPUT_DIR,
            num_train_epochs=NUM_EPOCHS,
            per_device_train_batch_size=BATCH_SIZE,
            gradient_accumulation_steps=GRAD_ACCUM,
            learning_rate=LEARNING_RATE,
            warmup_steps=WARMUP_STEPS,
            optim="adamw_8bit",
            fp16=not is_bfloat16_supported(),
            bf16=is_bfloat16_supported(),
            logging_steps=10,
            save_steps=200,
            save_total_limit=2,
            save_only_model=True,
            report_to="none",
        ),
        train_dataset=dataset,
        data_collator=data_collator,
    )

    trainer.train()

    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

if __name__ == "__main__":
    main()