ethical-hacking-llm-colab / INFERENCE_CELL.md
asdf98's picture
Add standalone inference notebook cell
e689e18 verified
|
Raw
History Blame Contribute Delete
4.64 kB

Standalone Inference Cell

Paste this as the last cell in Kaggle/Colab after training, or run it in a fresh notebook. It is independent of previous training cells.

Configure paths

Set:

BASE_MODEL_ID = "LiquidAI/LFM2.5-1.2B-Instruct"
ADAPTER_PATH = "./lfm25-stable-qlora-cybersecurity-adapter"

For Unsloth-trained adapters, examples:

BASE_MODEL_ID = "unsloth/LFM2.5-1.2B-Instruct"
ADAPTER_PATH = "./lfm25-lora-adapter"
BASE_MODEL_ID = "unsloth/Qwen3-4B-Instruct-2507-unsloth-bnb-4bit"
ADAPTER_PATH = "./qwen3-lora-adapter"

If you pushed adapter to Hub:

ADAPTER_PATH = "your-username/your-adapter-repo"

Complete independent cell

# ============================================================
# Standalone LoRA Adapter Inference / Chat Cell
# Works in fresh Kaggle/Colab notebook too.
# ============================================================
!pip install -q -U "transformers>=4.56.0" "peft>=0.18.0" "accelerate>=1.0.0" "bitsandbytes>=0.45.0" "huggingface_hub>=0.25.0"

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel

# -------------------- CONFIG --------------------
# Use the SAME base model used during training.
BASE_MODEL_ID = "LiquidAI/LFM2.5-1.2B-Instruct"

# Local adapter folder OR HF repo id.
# Common local paths:
#   ./lfm25-lora-adapter
#   ./qwen3-lora-adapter
#   ./gemma4-lora-adapter
#   ./lfm25-stable-qlora-cybersecurity-adapter
ADAPTER_PATH = "./lfm25-stable-qlora-cybersecurity-adapter"

LOAD_IN_4BIT = True
MAX_NEW_TOKENS = 512
TEMPERATURE = 0.7
TOP_P = 0.9
SYSTEM_PROMPT = "You are a helpful assistant. For cybersecurity topics, provide ethical, defensive, authorized guidance only."
# ------------------------------------------------

compute_dtype = torch.float16

bnb_config = None
if LOAD_IN_4BIT:
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_use_double_quant=True,
        bnb_4bit_compute_dtype=compute_dtype,
    )

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID, trust_remote_code=True, use_fast=True)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

print("Loading base model...")
base_model = AutoModelForCausalLM.from_pretrained(
    BASE_MODEL_ID,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True,
    low_cpu_mem_usage=True,
)

print("Loading LoRA adapter...")
model = PeftModel.from_pretrained(base_model, ADAPTER_PATH, is_trainable=False)
model.eval()

if torch.cuda.is_available():
    print(f"VRAM loaded: {torch.cuda.memory_allocated()/1e9:.2f} GB / {torch.cuda.get_device_properties(0).total_memory/1e9:.2f} GB")
print("✅ Model + adapter ready")


def build_prompt(messages):
    try:
        return tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
    except Exception:
        text = ""
        for m in messages:
            text += f"<{m['role']}>\n{m['content']}\n</{m['role']}>\n"
        text += "<assistant>\n"
        return text


def ask(prompt, history=None):
    if history is None:
        history = [{"role": "system", "content": SYSTEM_PROMPT}]
    messages = history + [{"role": "user", "content": prompt}]
    full_prompt = build_prompt(messages)
    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=MAX_NEW_TOKENS,
            do_sample=True,
            temperature=TEMPERATURE,
            top_p=TOP_P,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    new_tokens = output[0][inputs["input_ids"].shape[1]:]
    reply = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()
    return reply

# Single test prompt
reply = ask("Explain how parameterized queries prevent SQL injection with a safe Python example.")
print("Assistant:\n", reply)

Optional interactive chat loop

Run this after the cell above:

history = [{"role": "system", "content": SYSTEM_PROMPT}]

while True:
    user = input("You: ").strip()
    if user.lower() in {"exit", "quit", "q"}:
        break
    if not user:
        continue

    reply = ask(user, history)
    print("Assistant:", reply, "\n")
    history.append({"role": "user", "content": user})
    history.append({"role": "assistant", "content": reply})