import subprocess
import sys
from unsloth import FastLanguageModel
from trl import SFTTrainer
from transformers import TrainingArguments
from datasets import load_dataset

subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])

model, tokenizer = FastLanguageModel.from_pretrained(
    "unsloth/gemma-2-2b-it",
    max_seq_length = 2048,
    load_in_4bit = True,
)

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

dataset = load_dataset("json", data_files="python_security_dataset.json", split="train")

trainer = SFTTrainer(
    model = model,
    tokenizer = tokenizer,
    train_dataset = dataset,
    dataset_text_field = "messages",
    max_seq_length = 2048,
    args = TrainingArguments(
        per_device_train_batch_size = 2,
        gradient_accumulation_steps = 4,
        warmup_steps = 10,
        max_steps = 300,
        learning_rate = 2e-4,
        fp16 = True,
        logging_steps = 1,
        output_dir = "k1ng_final",
        optim = "adamw_8bit",
    ),
)

trainer.train()

model.save_pretrained("k1ng_by_alikay_h")
tokenizer.save_pretrained("k1ng_by_alikay_h")

# آپلود به HF
from huggingface_hub import notebook_login, HfApi
notebook_login()
api = HfApi()
api.upload_folder(folder_path="k1ng_by_alikay_h", repo_id="alikayh/k1ng-v1", repo_type="model")