{
 "cells": [
  {"cell_type":"markdown","metadata":{},"source":["# 🔐 Gemma-4 E2B Unsloth QLoRA — Root-Fixed Kaggle/Colab Notebook\n","\n","Keeps **Unsloth** for low VRAM and applies the same root fix for the `int.mean()` / `num_items_in_batch` bug.\n","\n","⚠️ Gemma-4 E2B is multimodal (`AutoModelForImageTextToText`) and tighter on free T4. Prefer LFM2.5 first. Defaults here are conservative.\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 1. Clean install Unsloth stack — run first\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["import os, sys, shutil, subprocess, pathlib\n","work_dir = pathlib.Path('/kaggle/working') if pathlib.Path('/kaggle/working').exists() else pathlib.Path('/content')\n","marker = work_dir / '.bex_unsloth_env_ready_v3'\n","shutil.rmtree(str(work_dir / 'unsloth_compiled_cache'), ignore_errors=True)\n","print('✅ Removed stale unsloth_compiled_cache')\n","if not marker.exists():\n","    print('Installing/updating Unsloth stack. Kernel will restart after this cell. Run all cells again after restart.')\n","    subprocess.check_call([sys.executable, '-m', 'pip', 'uninstall', '-y', 'unsloth', 'unsloth_zoo'])\n","    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '-U', '--no-cache-dir', 'unsloth', 'unsloth_zoo'])\n","    marker.write_text('ready')\n","    print('✅ Installed Unsloth. Restarting kernel now...')\n","    os.kill(os.getpid(), 9)\n","else:\n","    print('✅ Unsloth stack already prepared for this session')\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 2. Optional HF login\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["from huggingface_hub import login\n","# login(token='hf_YOUR_WRITE_TOKEN')\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 3. Imports and version check\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["from unsloth import FastLanguageModel, is_bfloat16_supported\n","from trl import SFTTrainer, SFTConfig\n","from datasets import load_dataset, concatenate_datasets\n","import torch, random, os, time\n","import transformers, trl, peft, accelerate\n","try:\n","    import unsloth\n","    print('unsloth', getattr(unsloth, '__version__', 'unknown'))\n","except Exception as e:\n","    print('unsloth version unavailable', e)\n","print('transformers', transformers.__version__)\n","print('trl', trl.__version__)\n","print('peft', peft.__version__)\n","print('accelerate', accelerate.__version__)\n","if torch.cuda.is_available():\n","    print('GPU:', torch.cuda.get_device_name(0))\n","    print(f'VRAM: {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB')\n","else:\n","    print('⚠️ No GPU found. Enable GPU runtime.')\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 4. Configuration\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["MODEL_ID = 'unsloth/gemma-4-E2B-it-unsloth-bnb-4bit'\n","RUN_NAME = 'gemma4-e2b-cyber-unsloth-rootfixed'\n","DATASET_CHOICE = 'cybersecurity'\n","SAMPLE_SIZE = 10000\n","MAX_SEQ_LENGTH = 1024\n","LORA_R = 8\n","LORA_ALPHA = 16\n","BATCH_SIZE = 1\n","GRAD_ACCUM = 8\n","MAX_STEPS = 1000\n","LEARNING_RATE = 1e-4\n","WARMUP_STEPS = 50\n","SAVE_STEPS = 250\n","LOGGING_STEPS = 10\n","PACKING = False\n","SEED = 3407\n","OUTPUT_DIR = './outputs_gemma4_rootfixed'\n","HUB_MODEL_ID = 'your-username/gemma4-e2b-cyber-unsloth-rootfixed'\n","PUSH_TO_HUB = False\n","random.seed(SEED)\n","torch.manual_seed(SEED)\n","if torch.cuda.is_available(): torch.cuda.manual_seed_all(SEED)\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 5. Load model with Unsloth 4-bit QLoRA\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["model, tokenizer = FastLanguageModel.from_pretrained(\n","    model_name=MODEL_ID,\n","    max_seq_length=MAX_SEQ_LENGTH,\n","    dtype=None,\n","    load_in_4bit=True,\n","    device_map={'': torch.cuda.current_device()} if torch.cuda.is_available() else None,\n",")\n","if tokenizer.pad_token is None:\n","    tokenizer.pad_token = tokenizer.eos_token\n","tokenizer.padding_side = 'right'\n","model = FastLanguageModel.get_peft_model(\n","    model,\n","    r=LORA_R,\n","    target_modules=['q_proj','k_proj','v_proj','o_proj','gate_proj','up_proj','down_proj'],\n","    lora_alpha=LORA_ALPHA,\n","    lora_dropout=0,\n","    bias='none',\n","    use_gradient_checkpointing='unsloth',\n","    random_state=SEED,\n","    use_rslora=True,\n","    loftq_config=None,\n",")\n","trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n","total = sum(p.numel() for p in model.parameters())\n","print(f'Trainable params: {trainable:,} / {total:,} ({100*trainable/total:.2f}%)')\n","if torch.cuda.is_available(): print(f'VRAM after model load: {torch.cuda.memory_allocated()/1e9:.2f} GB')\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 6. Load and format cybersecurity dataset\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["SYSTEM_SAFE = 'You are a cybersecurity education assistant. Provide defensive, ethical, and authorized security guidance only. Refuse harmful or unauthorized requests.'\n","def convert_sua(example):\n","    return {'messages': [{'role':'system','content':example.get('system') or SYSTEM_SAFE},{'role':'user','content':example['user']},{'role':'assistant','content':example['assistant']}]}\n","ds1=load_dataset('AlicanKiraz0/Cybersecurity-Dataset-Fenrir-v2.1', split='train').map(convert_sua, remove_columns=['system','user','assistant'])\n","ds2=load_dataset('Trendyol/Trendyol-Cybersecurity-Instruction-Tuning-Dataset', split='train').map(convert_sua, remove_columns=['system','user','assistant'])\n","dataset=concatenate_datasets([ds1,ds2])\n","print('Rows:', len(dataset))\n","if SAMPLE_SIZE and len(dataset)>SAMPLE_SIZE:\n","    dataset=dataset.shuffle(seed=SEED).select(range(SAMPLE_SIZE))\n","    print('Subsampled:', len(dataset))\n","print(dataset[0]['messages'])\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 7. Convert messages to text\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["def to_text(example):\n","    try:\n","        text=tokenizer.apply_chat_template(example['messages'], tokenize=False, add_generation_prompt=False)\n","    except Exception:\n","        text='\\n'.join([f\"<{m['role']}>\\n{m['content']}\\n</{m['role']}>\" for m in example['messages']]) + tokenizer.eos_token\n","    return {'text': text}\n","dataset=dataset.map(to_text, remove_columns=['messages'])\n","print(dataset[0]['text'][:500])\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 8. Train with current SFTConfig API\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["args = SFTConfig(\n","    output_dir=OUTPUT_DIR, dataset_text_field='text', max_length=MAX_SEQ_LENGTH, packing=PACKING,\n","    per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRAD_ACCUM,\n","    warmup_steps=WARMUP_STEPS, max_steps=MAX_STEPS, learning_rate=LEARNING_RATE,\n","    fp16=not is_bfloat16_supported(), bf16=is_bfloat16_supported(), logging_steps=LOGGING_STEPS,\n","    optim='adamw_8bit', weight_decay=0.01, lr_scheduler_type='linear', seed=SEED,\n","    save_strategy='steps', save_steps=SAVE_STEPS, save_total_limit=2, report_to='none',\n","    disable_tqdm=True, logging_first_step=True,\n",")\n","trainer = SFTTrainer(model=model, processing_class=tokenizer, train_dataset=dataset, args=args)\n","trainer.model_accepts_loss_kwargs = False\n","if hasattr(trainer, 'model'): trainer.model.config.use_cache = False\n","if torch.cuda.is_available(): print(f'VRAM before train: {torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB')\n","trainer_stats = trainer.train()\n","print(trainer_stats)\n"]},
  {"cell_type":"markdown","metadata":{},"source":["## 9. Save adapter and test\n"]},
  {"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":["model.save_pretrained('./gemma4-lora-adapter')\n","tokenizer.save_pretrained('./gemma4-lora-adapter')\n","print('Saved adapter to ./gemma4-lora-adapter')\n","if PUSH_TO_HUB:\n","    model.push_to_hub(HUB_MODEL_ID); tokenizer.push_to_hub(HUB_MODEL_ID)\n","FastLanguageModel.for_inference(model)\n","messages=[{'role':'system','content':SYSTEM_SAFE},{'role':'user','content':'Explain parameterized queries for preventing SQL injection with safe Python.'}]\n","inputs=tokenizer.apply_chat_template(messages, tokenize=True, add_generation_prompt=True, return_tensors='pt').to(model.device)\n","with torch.no_grad():\n","    outputs=model.generate(input_ids=inputs, max_new_tokens=256, temperature=0.7, top_p=0.9, do_sample=True, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id)\n","print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))\n"]}
 ],
 "metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10"}},
 "nbformat":4,
 "nbformat_minor":5
}