Update README.md

README.md

@@ -1,199 +1,749 @@
 ---
 library_name: transformers
-tags: []
+datasets:
+- nvidia/Aegis-AI-Content-Safety-Dataset-1.0
 ---
 
-# Model Card for Model ID
+# Model Card for AC/Meta-Llama-Guard-2-8B_Nvidia-Aegis-AI-Safety
 
 <!-- Provide a quick summary of what the model is/does. -->
-
-
-
-## Model Details
-
-### Model Description
-
-<!-- Provide a longer summary of what this model is. -->
-
-This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
-
-- **Developed by:** [More Information Needed]
-- **Funded by [optional]:** [More Information Needed]
-- **Shared by [optional]:** [More Information Needed]
-- **Model type:** [More Information Needed]
-- **Language(s) (NLP):** [More Information Needed]
-- **License:** [More Information Needed]
-- **Finetuned from model [optional]:** [More Information Needed]
-
-### Model Sources [optional]
-
-<!-- Provide the basic links for the model. -->
-
-- **Repository:** [More Information Needed]
-- **Paper [optional]:** [More Information Needed]
-- **Demo [optional]:** [More Information Needed]
-
-## Uses
-
-<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
-
-### Direct Use
-
-<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
-
-[More Information Needed]
-
-### Downstream Use [optional]
-
-<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
-
-[More Information Needed]
-
-### Out-of-Scope Use
-
-<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
-
-[More Information Needed]
-
-## Bias, Risks, and Limitations
-
-<!-- This section is meant to convey both technical and sociotechnical limitations. -->
-
-[More Information Needed]
-
-### Recommendations
-
-<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
-
-Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+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. 
+
+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
 
-Use the code below to get started with the model.
-
-[More Information Needed]
-
-## Training Details
-
-### Training Data
-
-<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
-
-[More Information Needed]
-
-### Training Procedure
-
-<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
-
-#### Preprocessing [optional]
-
-[More Information Needed]
-
-
-#### Training Hyperparameters
-
-- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
-
-#### Speeds, Sizes, Times [optional]
-
-<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
-
-[More Information Needed]
+```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.
+
+<BEGIN UNSAFE CONTENT 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
+<END UNSAFE CONTENT CATEGORIES>
+
+<BEGIN CONVERSATION>
+
+User: {user_prompt}
+
+<END CONVERSATION>
+
+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
-
-<!-- This section describes the evaluation protocols and provides the results. -->
-
-### Testing Data, Factors & Metrics
-
-#### Testing Data
-
-<!-- This should link to a Dataset Card if possible. -->
-
-[More Information Needed]
-
-#### Factors
-
-<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
-
-[More Information Needed]
-
-#### Metrics
-
-<!-- These are the evaluation metrics being used, ideally with a description of why. -->
-
-[More Information Needed]
-
-### Results
-
-[More Information Needed]
-
-#### Summary
-
-
-
-## Model Examination [optional]
-
-<!-- Relevant interpretability work for the model goes here -->
-
-[More Information Needed]
-
-## Environmental Impact
-
-<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
-
-Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
-
-- **Hardware Type:** [More Information Needed]
-- **Hours used:** [More Information Needed]
-- **Cloud Provider:** [More Information Needed]
-- **Compute Region:** [More Information Needed]
-- **Carbon Emitted:** [More Information Needed]
-
-## Technical Specifications [optional]
-
-### Model Architecture and Objective
-
-[More Information Needed]
-
-### Compute Infrastructure
-
-[More Information Needed]
-
-#### Hardware
-
-[More Information Needed]
-
-#### Software
-
-[More Information Needed]
-
-## Citation [optional]
-
-<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
-
-**BibTeX:**
-
-[More Information Needed]
-
-**APA:**
-
-[More Information Needed]
-
-## Glossary [optional]
-
-<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
-
-[More Information Needed]
-
-## More Information [optional]
-
-[More Information Needed]
-
-## Model Card Authors [optional]
-
-[More Information Needed]
-
-## Model Card Contact
-
-[More Information Needed]
+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.
+
+<BEGIN UNSAFE CONTENT 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
+<END UNSAFE CONTENT CATEGORIES>
+
+<BEGIN CONVERSATION>
+
+User: {user_prompt}
+
+<END CONVERSATION>
+
+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)
+```