This home assignement focuses on building a text classification model that can identify potentially harmful "prompt injection" attacks versus benign inputs. The model is based on a Convolutional Neural Network (CNN) architecture, specifically designed to capture key patterns in text. This choice of a TextCNN model is motivated by its effectiveness in natural language processing tasks where the structure and patterns of text sequences are crucial. The CNN model can efficiently learn relevant features from the text with its convolutional layers, which process data in chunks (subsequences), making it well-suited for extracting local dependencies in text.

For this model, each text prompt is converted into a feature vector where each character in the prompt is represented by its ASCII value, truncated or padded to a fixed length of 256 characters. This preprocessing approach keeps the model input simple and lightweight, and the fixed-length feature vectors make it easier to batch-process data.

The model consists of an embedding layer followed by multiple 1D convolutional layers with varying kernel sizes. The embedding layer converts each character into a dense vector representation, while the convolutional layers use different kernel sizes to capture diverse patterns in the text (such as specific sequences that may signal benign or harmful intent). These features are then passed through a fully connected (linear) layer to classify each prompt as "benign" or "prompt injection." We use Cross-Entropy Loss to optimize the binary classification, and the Adam optimizer to enable faster convergence.

During training, the dataset is divided into training, validation, and test sets to monitor the model’s performance and prevent overfitting. The model is trained over a few epochs, with each epoch reporting training and validation loss, as well as additional performance metrics, including accuracy, precision, recall, and F1 score on the validation set. The model saves its best version based on validation loss, which helps in preserving a model that generalizes well on unseen data. After training, the best model is evaluated on the test set, and a confusion matrix is generated to visualize the results, highlighting areas where the model might misclassify benign prompts as harmful or vice versa.

The project includes a deployment function to make it easy to use the trained model for real-time prompt classification. With the function classify_prompt, users can input new prompts to receive a classification label and a confidence score. This setup allows for a practical and efficient inference pipeline, suitable for applications where text-based threat detection is needed.

To run the code, install the required Python libraries (torch, pandas, numpy, matplotlib, sklearn, and huggingface_hub). The script downloads the training data from Hugging Face Hub and trains the model, saving the best model weights to a file named best_model.pth. This file can then be loaded to classify new text prompts using the provided classify_prompt function. Overall, this project offers a streamlined solution for text classification, with a focus on detecting harmful prompt patterns in a simple, interpretable, and reusable way.