4-bit NF4 quantization for llama guard 3

this quantization will help for users with low vram and get desired result but note using 4 bit quantization has its drawback.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

MODEL_ID = "meta-llama/Llama-Guard-3-8B"

# Configure 4-bit quantization to drastically reduce VRAM usage
quantization_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",           # Use NormalFloat4 quantization
    bnb_4bit_compute_dtype=torch.bfloat16, # Perform calculations in bfloat16 for speed
    bnb_4bit_use_double_quant=True      # Double quantization saves extra memory
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

# Load model using the 4-bit config
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    quantization_config=quantization_config,
    device_map="auto"
)

def moderate_input(user_prompt: str) -> str:
    """
    Checks the user's incoming prompt.
    Returns 'safe' or 'unsafe \n [Category Code]'.
    """
    chat = [
        {"role": "user", "content": user_prompt}
    ]
    
    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(model.device)
    prompt_len = input_ids.shape[-1]
    
    with torch.no_grad():
        output = model.generate(input_ids, max_new_tokens=10, pad_token_id=0)
    
    return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True).strip()


def moderate_output(user_prompt: str, ai_response: str) -> str:
    """
    Checks the AI's generated response against the user's original prompt.
    Returns 'safe' or 'unsafe \n [Category Code]'.
    """
    chat = [
        {"role": "user", "content": user_prompt},
        {"role": "assistant", "content": ai_response}
    ]
    
    input_ids = tokenizer.apply_chat_template(chat, return_tensors="pt").to(model.device)
    prompt_len = input_ids.shape[-1]
    
    with torch.no_grad():
        output = model.generate(input_ids, max_new_tokens=10, pad_token_id=0)
    
    return tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True).strip()