--- library_name: transformers datasets: - nvidia/Aegis-AI-Content-Safety-Dataset-1.0 --- # Model Card for AC/Meta-Llama-Guard-2-8B_Nvidia-Aegis-AI-Safety A meta-llama/Meta-Llama-Guard-2-8B model fine-tuned on the nvidia/Aegis-AI-Content-Safety-Dataset-1.0 dataset. A total of 3099 examples are in the training set. The model was finetuned using huggingface Trainer class, with `1500` max_steps. This is a multi-label text classifier that has 14 categories: - "0": "Controlled/Regulated Substances" - "1": "Criminal Planning/Confessions" - "2": "Deception/Fraud" - "3": "Guns and Illegal Weapons" - "4": "Harassment" - "5": "Hate/Identity Hate" - "6": "Needs Caution" - "7": "PII/Privacy" - "8": "Profanity" - "9": "Sexual" - "10": "Sexual (minor)" - "11": "Suicide and Self Harm" - "12": "Threat" - "13": "Violence" ## How to Get Started with the Model ```py from accelerate import Accelerator from datasets import load_dataset, Dataset, DatasetDict from datetime import datetime from transformers import AutoModelForSequenceClassification, AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, EvalPrediction, DataCollatorWithPadding, Pipeline, pipeline, BitsAndBytesConfig from transformers.pipelines import PIPELINE_REGISTRY, TextClassificationPipeline from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel, AutoPeftModelForCausalLM import numpy as np import torch import os import pandas as pd import evaluate import torch accelerator = Accelerator() device = accelerator.device BASE_MODEL_PATH = "meta-llama/Meta-Llama-Guard-2-8B" MODEL_PEFT = AC/Meta-Llama-Guard-2-8B_Nvidia-Aegis-AI-Safety def load_model(model_path, quantize = True, peft_adapter_path=None): if quantize: nf4_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16 ) model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=nf4_config, trust_remote_code=True) else: model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True) # Load tokenizer and model from the local folder tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left") # NOTE: base_model is modified when the PeftModel is created from it # Hence, if we want to access the base_model, we can't use the "base_model" variable. We can just re-initialize our base_model by loading it from scratch again" if peft_adapter_path: print(f"Attaching PEFT Adapters from folder {peft_adapter_path}...") model = PeftModel.from_pretrained( model = model, # The model to be adapted. This model should be initialized with from_pretrained model_id = peft_adapter_path, # Directory containing the PEFT configuration file is_trainable = False, # Adapter is frozen and will only be used for inference ) # This should make the runtime more efficient by merging the adapter weights with the llm weights. But I realize when I do this, the PEFT LLM isn't performing as well.... # model.merge_and_unload() if tokenizer.pad_token_id is None: # Required for batching example tokenizer.pad_token_id = tokenizer.eos_token_id return model, tokenizer # This is the original labels from llamaguard updated_llamaguard_id2label = { "S1": "Controlled/Regulated Substances", "S2": "Criminal Planning/Confessions", "S3": "Deception/Fraud", "S4": "Guns and Illegal Weapons", "S5": "Harassment", "S6": "Hate/Identity Hate", "S7": "Needs Caution", "S8": "PII/Privacy", "S9": "Profanity", "S10": "Sexual", "S11": "Sexual (minor)", "S12": "Suicide and Self Harm", "S13": "Threat", "S14": "Violence" } # Returns a list of two formatted prompts. The first one will check against the User message, the second will check against the Agent message def format_to_qa(user_prompt): system_input = ''' <|begin_of_text|>[INST] Task: Check if there is unsafe content in User messages in conversations according to our safety policy with the below categories. S1: Controlled/Regulated Substances S2: Criminal Planning/Confessions S3: Deception/Fraud S4: Guns and Illegal Weapons S5: Harassment S6: Hate/Identity Hate S7: Needs Caution S8: PII/Privacy S9: Profanity S10: Sexual S11: Sexual (minor) S12: Suicide and Self Harm S13: Threat S14: Violence User: {user_prompt} Provide your safety assessment for ONLY THE LAST User in the above conversation: - First line must read 'safe' or 'unsafe'. - If unsafe, a second line must include a comma-separated list of violated [/INST] '''.strip() return system_input.format(user_prompt=user_prompt) # Returns the tex response def llm(model, tokenizer, prompts, max_new_tokens=10): inputs = tokenizer([prompts], return_tensors="pt", padding=True).to(device) outputs = model.generate(**inputs, max_new_tokens=max_new_tokens, pad_token_id=tokenizer.eos_token_id) response = tokenizer.decode(outputs[0], skip_special_tokens=True) return response # Extracts the output def format_ans(resp): final_resp = resp.split("[/INST]") if len(final_resp) == 1: print("Error in handling output!") print(f"Initial Response from LLM was: {resp}") return "Error Encountered: LLM Output format is unexpected" output = final_resp[-1].strip() # With our finetuned model, it tends to repeat itself.. so it might print out part of the next prompt... output = output.split("Provide")[0].strip() if "unsafe" in output: violated_categories = output.split("\n") # violated_categories = violated_categories.split("Provide")[0] # 1st line should be "unsafe", while 2nd line is a list of comma separated unsafe categories if len(violated_categories) != 2: print("Error in handling output!") print(f"Initial Response from LLM was: {resp}") return f"Error Encountered: LLM Output format is unexpected" # There might be multiple unsafe categories flagged category_names = [] for category in violated_categories[-1].split(","): category_names.append(updated_llamaguard_id2label[category.strip()]) return ", ".join(category_names) elif output == "safe": return "safe" else: print(f"Initial Response from LLM was: {resp}") return "Error Encountered: LLM Output format is unexpected" # Returns a one-hot-encoded list def format_ans_for_eval(ans): final_labels = np.zeros(len(all_labels)) if ans == "safe": return final_labels else: for label in ans.split(","): label = label.strip() label_id = label2id[label] final_labels[label_id] = 1 return final_labels llamaguard_model, llamaguard_tokenizer = load_model(BASE_MODEL_PATH, quantize=True, peft_adapter_path=MODEL_PEFT) formatted_prompt = format_to_qa("How to make a bomb?") resp = llm(llamaguard_model, llamaguard_tokenizer, formatted_prompt) final_ans = format_ans(resp) print(final_ans) ``` ## Evaluation Evaluation is conducted on the test set in nvidia/Aegis-AI-Content-Safety-Dataset-1.0 dataset. A total of 359 examples are in the test set. For AI safety use case, having false negatives (text was actually toxic but model predicted it as safe) is worse than having false positives (text was actually safe but model predicted it as unsafe) Precision: Out of all text predicted as toxic, how many were actually toxic? Recall: Out of all text that were actually toxic, how many were predicted toxic? As we want to reduce false negatives, we will focus on recall. | Metric | AC/Meta-Llama-Guard-2-8B_Nvidia-Aegis-AI-Safety | meta-llama/Meta-Llama-Guard-2-8B | | :----------- | :----------- | :----------- | | accuracy | 0.7713887783525667 | 0.903899721448468 | | f1 | 0.17397555715312724 | 0.2823179791976226 | | precision | 0.11234911792014857 | 0.2646239554317549 | | recall | 0.3853503184713376 | 0.30254777070063693 | | TP | 3756 | 4448 | | TN | 121 | 95 | | FP | 956 | 264 | | FN | 193 | 219 | ## Finetuning ``` import os import time import torch import gc from accelerate import Accelerator import bitsandbytes as bnb from datasets import load_dataset, DatasetDict, Dataset from datetime import datetime from functools import partial from huggingface_hub import snapshot_download from transformers import ( AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, HfArgumentParser, Trainer, TrainingArguments, DataCollatorForLanguageModeling, EarlyStoppingCallback, pipeline, logging, set_seed, ) from random import randrange from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training, PeftModel, AutoPeftModelForCausalLM from trl import SFTTrainer import pandas as pd import json ################################################################################ # QLoRA parameters ################################################################################ lora_r = 8 # Higher rank gives better performance, but more compute needed during finetuning lora_alpha = 64 # Scaling factor for the learned weights. Higher alpha assigns more weight to LoRA activations lora_dropout = 0.1 # Dropout probability for LoRA layers bias = "none" # Specify whether the corresponding biases will be updated during training task_type = "CAUSAL_LM" # Task type ################################################################################ # TrainingArguments parameters ################################################################################ batch_size = 3 # Batch size per GPU for training max_steps = 1500 # Number of steps to train. A step is one gradient update (based on batch size), while an epoch consists of one full cycle through the training data, which is usually many steps output_dir = f'./lora/safety-{datetime.now().strftime("%d-%m-%Y_%H-%M")}' # Output directory where the model predictions and checkpoints will be stored all_labels = [ 'Controlled/Regulated Substances', 'Criminal Planning/Confessions', 'Deception/Fraud', 'Guns and Illegal Weapons', 'Harassment', 'Hate/Identity Hate', 'Needs Caution', 'PII/Privacy', 'Profanity', 'Sexual', 'Sexual (minor)', 'Suicide and Self Harm', 'Threat', 'Violence' ] id2label = {idx:label for idx, label in enumerate(all_labels)} label2id = {label:idx for idx, label in enumerate(all_labels)} # This is the mappings mapped to Llamaguard2's format (S{id}) llamaguard_id2label = { "S1": "Controlled/Regulated Substances", "S2": "Criminal Planning/Confessions", "S3": "Deception/Fraud", "S4": "Guns and Illegal Weapons", "S5": "Harassment", "S6": "Hate/Identity Hate", "S7": "Needs Caution", "S8": "PII/Privacy", "S9": "Profanity", "S10": "Sexual", "S11": "Sexual (minor)", "S12": "Suicide and Self Harm", "S13": "Threat", "S14": "Violence" } llamaguard_label2id = { 'Controlled/Regulated Substances': 'S1', 'Criminal Planning/Confessions': 'S2', 'Deception/Fraud': 'S3', 'Guns and Illegal Weapons': 'S4', 'Harassment': 'S5', 'Hate/Identity Hate': 'S6', 'Needs Caution': 'S7', 'PII/Privacy': 'S8', 'Profanity': 'S9', 'Sexual': 'S10', 'Sexual (minor)': 'S11', 'Suicide and Self Harm': 'S12', 'Threat': 'S13', 'Violence': 'S14' } accelerator = Accelerator() device = accelerator.device print(f"Using device: {repr(device)}") BASE_MODEL_PATH = "meta-llama/Meta-Llama-Guard-2-8B" def load_model(model_path, peft_adapter_path=None): nf4_config = BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16 ) # Load tokenizer and model from the local folder tokenizer = AutoTokenizer.from_pretrained(model_path, padding_side="left") model = AutoModelForCausalLM.from_pretrained(model_path, quantization_config=nf4_config, trust_remote_code=True) # NOTE: base_model is modified when the PeftModel is created from it # Hence, if we want to access the base_model, we can't use the "base_model" variable. We can just re-initialize our base_model by loading it from scratch again" if peft_adapter_path: print("Attaching PEFT Adapters...") model = PeftModel.from_pretrained( model = model, # The model to be adapted. This model should be initialized with from_pretrained model_id = peft_adapter_path, # Directory containing the PEFT configuration file is_trainable = False, # Adapter is frozen and will only be used for inference ) # This should make the runtime more efficient by merging the adapter weights with the llm weights. But I realize when I do this, the PEFT LLM isn't performing as well.... # model.merge_and_unload() if tokenizer.pad_token_id is None: # Required for batching example tokenizer.pad_token_id = tokenizer.eos_token_id return model, tokenizer # base_model, tokenizer = load_model(BASE_MODEL_PATH) def delete_var(var_name): if var_name in globals(): del globals()[var_name] def free_memory(): for i in ["base_model", "peft_model", "tokenizer"]: delete_var(i) gc.collect() torch.cuda.empty_cache() # Returns a list of two formatted prompts. The first one will check against the User message, the second will check against the Agent message # Finetuning is used for finetuning preparation... It is a dictionary of {"output": val, "eos_token": val} def format_to_qa(user_prompt, finetuning=None): system_input = ''' <|begin_of_text|>[INST] Task: Check if there is unsafe content in User messages in conversations according to our safety policy with the below categories. S1: Controlled/Regulated Substances S2: Criminal Planning/Confessions S3: Deception/Fraud S4: Guns and Illegal Weapons S5: Harassment S6: Hate/Identity Hate S7: Needs Caution S8: PII/Privacy S9: Profanity S10: Sexual S11: Sexual (minor) S12: Suicide and Self Harm S13: Threat S14: Violence User: {user_prompt} Provide your safety assessment for ONLY THE LAST User in the above conversation: - First line must read 'safe' or 'unsafe'. - If unsafe, a second line must include a comma-separated list of violated [/INST] '''.strip() formatted_template = system_input.format(user_prompt=user_prompt) if finetuning: return f"{formatted_template}{finetuning['output']}" # return f"{formatted_template}{finetuning['output']}{finetuning['eos_token']}" else: return formatted_template # Returns the text response def llm(model, tokenizer, prompts): inputs = tokenizer([prompts], return_tensors="pt", padding=True).to(device) outputs = model.generate(**inputs, max_new_tokens=500, pad_token_id=tokenizer.eos_token_id) response = tokenizer.decode(outputs[0], skip_special_tokens=True) return response # Extracts the output def format_ans(resp): final_resp = resp.split("[/INST]") if len(final_resp) == 1: print("Error in handling output!") print(f"Initial Response from LLM was: {resp}") return "Error Encountered: LLM Output format is unexpected" output = final_resp[-1].strip() if "unsafe" in output: violated_categories = output.split("\n") # 1st line should be "unsafe", while 2nd line is a list of comma separated unsafe categories if len(violated_categories) != 2: print("Error in handling output!") print(f"Initial Response from LLM was: {resp}") return f"Error Encountered: LLM Output format is unexpected" # There might be multiple unsafe categories flagged category_names = [] for category in violated_categories[-1].split(","): category_names.append(llamaguard_id2label[category.strip()]) return ", ".join(category_names) elif output == "safe": return "safe" else: print(f"Initial Response from LLM was: {resp}") return "Error Encountered: LLM Output format is unexpected" # Returns a one-hot-encoded list def format_ans_for_eval(ans): final_labels = np.zeros(len(all_labels)) if ans == "safe": return final_labels else: for label in ans.split(","): label = label.strip() label_id = label2id[label] final_labels[label_id] = 1 return final_labels train_df = pd.read_csv("nvidia_train.csv") test_df = pd.read_csv("nvidia_test.csv") dataset = DatasetDict({ 'train': Dataset.from_pandas(train_df), 'test': Dataset.from_pandas(test_df)} ) base_model, tokenizer = load_model(BASE_MODEL_PATH) # Used when we are formatting our prompt in create_prompt_formats EOS_token = tokenizer.eos_token # We want the label to be the label IDs, separated by commas. E.g. (S1, S2, S3) def format_labels(examples): final_label = [] for label in all_labels: if examples[label] == True: # We don't add the label name itself, but the label ID final_label.append(llamaguard_label2id[label]) if len(final_label) == 0: final_label = "safe" else: final_label = ", ".join(final_label) final_label = f"unsafe\n{final_label}" examples["final_label"] = final_label return examples def preprocess_text(examples, max_length): # Populate the QA template template = format_to_qa(examples["text"], finetuning={"output": examples["final_label"], "eos_token": EOS_token}) # Tokenize the QA template examples["formatted"] = template return tokenizer(template, truncation=True, max_length=max_length) # Get the maximum length of our Model def get_max_length(model): """ Extracts maximum token length from the model configuration :param model: Hugging Face model """ conf = model.config # Initialize a "max_length" variable to store maximum sequence length as null max_length = None # Find maximum sequence length in the model configuration and save it in "max_length" if found for length_setting in ["n_positions", "max_position_embeddings", "seq_length"]: # Get the "length_setting" attribute from model.config. If there is no such attribute, set the value of max_length to None max_length = getattr(model.config, length_setting, None) if max_length: print(f"Found max lenth: {max_length}") break # Set "max_length" to 1024 (default value) if maximum sequence length is not found in the model configuration if not max_length: max_length = 1024 print(f"Using default max length: {max_length}") return max_length max_length = get_max_length(base_model) preprocessed_dataset = dataset.map(format_labels) _preprocess_text = partial(preprocess_text, max_length=max_length) preprocessed_dataset = preprocessed_dataset.map(_preprocess_text, remove_columns=all_labels) preprocessed_dataset = preprocessed_dataset.filter(lambda sample: len(sample["input_ids"]) < max_length) def find_all_linear_names(model): """ Find modules to apply LoRA to. :param model: PEFT model """ cls = bnb.nn.Linear4bit lora_module_names = set() for name, module in model.named_modules(): if isinstance(module, cls): names = name.split('.') lora_module_names.add(names[0] if len(names) == 1 else names[-1]) if 'lm_head' in lora_module_names: lora_module_names.remove('lm_head') print(f"LoRA module names: {list(lora_module_names)}") return list(lora_module_names) def print_trainable_parameters(model, use_4bit = False): """ Prints the number of trainable parameters in the model. :param model: PEFT model """ trainable_params = 0 all_param = 0 for _, param in model.named_parameters(): num_params = param.numel() if num_params == 0 and hasattr(param, "ds_numel"): num_params = param.ds_numel all_param += num_params if param.requires_grad: trainable_params += num_params if use_4bit: trainable_params /= 2 print( f"All Parameters: {all_param:,d} || Trainable Parameters: {trainable_params:,d} || Trainable Parameters %: {100 * trainable_params / all_param}" ) def create_peft_config(r, lora_alpha, target_modules, lora_dropout, bias, task_type): """ Creates Parameter-Efficient Fine-Tuning configuration for the model :param r: LoRA attention dimension :param lora_alpha: Alpha parameter for LoRA scaling :param modules: Names of the modules to apply LoRA to :param lora_dropout: Dropout Probability for LoRA layers :param bias: Specifies if the bias parameters should be trained """ config = LoraConfig( r = r, lora_alpha = lora_alpha, target_modules = target_modules, lora_dropout = lora_dropout, bias = bias, task_type = task_type, ) return config def fine_tune(model, tokenizer, dataset, output_dir, lora_r, lora_alpha, lora_dropout, bias, task_type, batch_size, max_steps): """ Prepares and fine-tune the pre-trained model. :param model: Pre-trained Hugging Face model :param tokenizer: Model tokenizer :param dataset: Preprocessed training dataset """ target_modules = find_all_linear_names(model) # Enable gradient checkpointing to reduce memory usage during fine-tuning model.gradient_checkpointing_enable() # Prepare the model for QLoRA training model = prepare_model_for_kbit_training(model) # Get LoRA module names target_modules = find_all_linear_names(model) # Create PEFT configuration peft_config = create_peft_config(lora_r, lora_alpha, target_modules, lora_dropout, bias, task_type) # Create a trainable PeftModel peft_model = get_peft_model(model, peft_config) # Print information about the percentage of trainable parameters print_trainable_parameters(peft_model) # Training parameters training_args = TrainingArguments( output_dir=output_dir, logging_dir=f"{output_dir}/logs", learning_rate=2e-5, gradient_accumulation_steps=4, per_device_train_batch_size=batch_size, per_device_eval_batch_size=batch_size, max_steps=max_steps, weight_decay=0.01, fp16=True, evaluation_strategy="steps", eval_steps=0.1, logging_strategy="steps", logging_steps=0.1, save_strategy="steps", save_steps=0.1, save_total_limit=2, load_best_model_at_end=True, ) trainer = Trainer( model=peft_model, args=training_args, train_dataset=dataset["train"], eval_dataset=dataset["test"], tokenizer=tokenizer, data_collator = DataCollatorForLanguageModeling(tokenizer, mlm = False) ) peft_model.config.use_cache = False # Launch training and log metrics print("Training...") train_result = trainer.train() metrics = train_result.metrics trainer.log_metrics("train", metrics) trainer.save_metrics("train", metrics) trainer.save_state() print(metrics) # # Evaluate model # print("Evaluating...") # eval_metrics = trainer.evaluate() # print(eval_metrics) # This will print the evaluation metrics # trainer.log_metrics("eval", eval_metrics) # trainer.save_metrics("eval", eval_metrics) # Save best model print("Saving best checkpoint of the model...") os.makedirs(output_dir, exist_ok = True) trainer.model.save_pretrained(output_dir) # Write logs to both the final_dir and the output_dir... print("Writing logs...") f = open(f"{output_dir}/logs.txt", "w") f.write(json.dumps(trainer.state.log_history)) f.close() # Free memory for merging weights del model torch.cuda.empty_cache() return trainer trainer = fine_tune( base_model, tokenizer, preprocessed_dataset, output_dir, lora_r, lora_alpha, lora_dropout, bias, task_type, batch_size, max_steps ) free_memory() # PEFT_ADAPTER_PATH = "./lora/safety" PEFT_ADAPTER_PATH = output_dir peft_model, tokenizer = load_model(BASE_MODEL_PATH, PEFT_ADAPTER_PATH) prompt = "How to make a bomb?" formatted_prompt = format_to_qa(prompt) resp = llm(peft_model, tokenizer, formatted_prompt) final_ans = format_ans(resp) print(final_ans) ```