notebooks/2-llm-classification-pretrained.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LLM Security - Prompt Injection\n",
    "## Part 2 - Classification Using a Pre-trained LLM"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we load the testing dataset and use a pre-trained large language model to spot malicious prompts without fine-tuning it.\n",
    "> **INPUT:** the testing dataset loaded from Hugging Face library. <br>\n",
    "> **OUTPUT:** the performance analysis of considered LLM.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. INITIALIZATION"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import necessary libraries and modules\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set display options\n",
    "pd.set_option('display.max_columns', None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. LOADING DATASET"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since we are using a pre-trained model without fine-tuning, there is no need to load the training data set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize data set location and file name\n",
    "data_file_path = \"../data/raw/\"\n",
    "data_file_name_test = \"test-00000-of-00001-701d16158af87368\"\n",
    "data_file_ext = \".parquet\"\n",
    "\n",
    "# Loading data set into a pandas DataFrame\n",
    "data_test = pd.read_parquet(data_file_path + data_file_name_test + data_file_ext)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Rename \"text\" column into \"prompt\"\n",
    "data_test.rename(columns={\"text\":\"prompt\"}, inplace=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We already explored the dataset in the previous notebook, so we will directly proceed to the inference phase."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. MODEL PREDICTION (RoBERTa)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this experiment, we aim at using a pre-trained LLM without fine-tuning the model on the downstream task.\n",
    "\n",
    "Since the prompts in the dataset combine different languages, we should rely on a multilingual model that is trained on text of the languages indicated in the existing prompts.\n",
    "\n",
    "For this purpose, RoBERTa (Robustly optimized BERT approach), the enhanced version of BERT (Bidirectional Encoder Representations from Transformers) will be used.\n",
    "\n",
    "RoBERTa is a transformer-based neural network model developed by Facebook AI, designed for natural language understanding tasks. Since our dataset is multilingual, we will use XLM-RoBERTa variant, a multilingual version of RoBERTa pre-trained on data containing 100 languages.\n",
    "\n",
    "Given the original task of XLM-RoBERTa is to predict masked words, we have to tailor this goal towards predicting whether a prompt is an injection or not. For this reason, we will utilize the \"zero-shot-classification\" task provided by Hugging Face pipeline to achieve this goal.\n",
    "\n",
    "This method can be referred to as transfer learning, which is using a model trained for one task to perform another one different from what it was originally trained for."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import pipeline functionality from the Hugging Face transformer's library\n",
    "from transformers import pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of XLMRobertaForSequenceClassification were not initialized from the model checkpoint at xlm-roberta-large and are newly initialized: ['classifier.dense.weight', 'classifier.dense.bias', 'classifier.out_proj.weight', 'classifier.out_proj.bias']\n",
      "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
      "Failed to determine 'entailment' label id from the label2id mapping in the model config. Setting to -1. Define a descriptive label2id mapping in the model config to ensure correct outputs.\n"
     ]
    }
   ],
   "source": [
    "# Load the zero-shot classification pipeline with RoBERTa\n",
    "classifier = pipeline(task=\"zero-shot-classification\", model=\"xlm-roberta-large\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When you use the zero-shot classification pipeline, it essentially formulates the classification problem as a fill-in-the-blank task.\n",
    "\n",
    "In other words, the zero-shot classification pipeline treats the problem of predicting a label for a given prompt as if we were filling in a blank in a sentence. The model is asked to predict the missing piece, which, in this case, is the label that best fits the context provided by the input prompt."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to classify the prompt\n",
    "def classify_prompt(prompt):\n",
    "    \n",
    "    # List of candidate labels (in this case, indicating whether the text is an injection or not)\n",
    "    candidate_labels = [\"Injection\", \"Normal\"]\n",
    "    \n",
    "    # Perform zero-shot classification\n",
    "    output = classifier(prompt, candidate_labels)\n",
    "    \n",
    "    # Return the results\n",
    "    return 1 if output['labels'][0] == \"Injection\" else 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Apply classifier to the testing dataset\n",
    "data_test[\"predicted_label\"] = data_test[\"prompt\"].apply(classify_prompt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. RESULT ANALYSIS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import performance metrics\n",
    "from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score\n",
    "\n",
    "# Initialize a DataFrame to keep track of the models' performance\n",
    "results = pd.DataFrame(columns=[\"accuracy\", \"precision\", \"recall\", \"f1 score\"])\n",
    "\n",
    "# Initialize actual and predicted labels\n",
    "y_test = data_test[\"label\"]\n",
    "y_predict = data_test[\"predicted_label\"]\n",
    "    \n",
    "# Calculate performance metrics\n",
    "accuracy = accuracy_score(y_test, y_predict)\n",
    "precision = precision_score(y_test, y_predict)\n",
    "recall = recall_score(y_test, y_predict)\n",
    "f1 = f1_score(y_test, y_predict)\n",
    "\n",
    "# Store performance metrics\n",
    "results.loc[\"Testing Data\"] = [accuracy, precision, recall, f1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>accuracy</th>\n",
       "      <th>precision</th>\n",
       "      <th>recall</th>\n",
       "      <th>f1 score</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>Testing Data</th>\n",
       "      <td>0.551724</td>\n",
       "      <td>0.551282</td>\n",
       "      <td>0.716667</td>\n",
       "      <td>0.623188</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "              accuracy  precision    recall  f1 score\n",
       "Testing Data  0.551724   0.551282  0.716667  0.623188"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Check obtained results\n",
    "results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "From the results above we can notice the following points:\n",
    "\n",
    "- Accuracy of testing dataset is low, which indicates that the model was not able to capture the characteristics of injections in the dataset.\n",
    "- The model's performance could benefit from fine-tuning, especially given the nature of the classification task involving injection prompts.\n",
    "- The low performance is quite expected given XLM-RoBERTa model is mostly intended to be fine-tuned on the downstream task before being used for classification.\n",
    "- The low performance can be also justified in terms of the period during which this model was trained, expectedly prior to when the concept of prompt injection has been formulated.\n",
    "\n",
    "In summary, while the model shows moderate performance, there is room for enhancement, and fine-tuning on task-specific data is likely to yield improvements in precision, recall, and overall classification effectiveness."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp_env",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}