Update README.md

README.md

@@ -1,53 +1,35 @@
----
-library_name: transformers
-license: apache-2.0
-base_model: google/electra-base-discriminator
-tags:
-- generated_from_trainer
-model-index:
-- name: electra-trainer
-  results: []
----
+# ELECTRA Trainer for Prompt Injection Detection
 
-<!-- This model card has been generated automatically according to the information the Trainer had access to. You
-should probably proofread and complete it, then remove this comment. -->
+ Google Colab notebook link: https://colab.research.google.com/drive/11da3m_gYwmkURcjGn8_kp23GiM-INDrm?usp=sharing
+## Overview
 
-# electra-trainer
+This repository contains a fine-tuned ELECTRA model designed for detecting prompt injections in AI systems. The model classifies input prompts into two categories: benign and jailbreak. This approach aims to enhance the safety and robustness of AI applications.
 
-This model is a fine-tuned version of [google/electra-base-discriminator](https://huggingface.co/google/electra-base-discriminator) on the None dataset.
+## Approach and Design Decisions
 
-## Model description
+The primary goal of this project was to create a reliable model that can distinguish between safe and potentially harmful prompts. Key design decisions included:
 
-More information needed
+- **Model Selection**: We chose the ELECTRA model due to its efficient training process and strong performance on text classification tasks. ELECTRA's architecture allows for effective learning from limited data, which is crucial given the specificity of our task.
+  
+- **Data Preparation**: A custom dataset was curated, consisting of diverse prompts labeled as either benign or jailbreak. The dataset aimed to balance both classes to mitigate biases during training.
 
-## Intended uses & limitations
 
-More information needed
+## Model Architecture and Training Strategy
 
-## Training and evaluation data
+The model is based on the `google/electra-base-discriminator` architecture. Here’s an overview of the training strategy:
 
-More information needed
+1. **Tokenization**: We utilized the ELECTRA tokenizer to prepare input prompts. Padding and truncation were handled to ensure uniform input sizes.
 
-## Training procedure
+2. **Training Configuration**:
+   - **Learning Rate**: Set to 5e-05 for stable convergence.
+   - **Batch Size**: A batch size of 16 was chosen to balance training speed and memory usage.
+   - **Epochs**: The model was trained for 2 epochs to prevent overfitting while still allowing sufficient learning from the dataset.
 
-### Training hyperparameters
+3. **Evaluation**: The model’s performance was evaluated on a validation set, focusing on metrics such as accuracy, precision, recall, and F1 score. 
 
-The following hyperparameters were used during training:
-- learning_rate: 5e-05
-- train_batch_size: 16
-- eval_batch_size: 16
-- seed: 42
-- optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
-- lr_scheduler_type: linear
-- num_epochs: 2
+## Key Results and Observations
 
-### Training results
+- The model achieved a high accuracy rate on the validation set, indicating its effectiveness in distinguishing between benign and harmful prompts.
 
-
-
-### Framework versions
-
-- Transformers 4.44.2
-- Pytorch 2.5.0+cu121
-- Datasets 3.1.0
-- Tokenizers 0.19.1
+## Instructions for Running the Inference Pipeline
+- Simply input your text into Classify_Prompt function and get your results as the classification itself and the confidence of the model.