add: docs for train classifier

docs/train_classifier.md

@@ -0,0 +1,3 @@
+# Train Classifier
+
+::: guardrails_genie.train_classifier

guardrails_genie/train_classifier.py

@@ -16,6 +16,22 @@ import wandb
 
 
 class StreamlitProgressbarCallback(TrainerCallback):
+    """
+    StreamlitProgressbarCallback is a custom callback for the Hugging Face Trainer
+    that integrates a progress bar into a Streamlit application. This class updates
+    the progress bar at each training step, providing real-time feedback on the
+    training process within the Streamlit interface.
+
+    Attributes:
+        progress_bar (streamlit.delta_generator.DeltaGenerator): A Streamlit progress
+            bar object initialized to 0 with the text "Training".
+
+    Methods:
+        on_step_begin(args, state, control, **kwargs):
+            Updates the progress bar at the beginning of each training step. The progress
+            is calculated as the percentage of completed steps out of the total steps.
+            The progress bar text is updated to show the current step and the total steps.
+    """
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -42,6 +58,8 @@ def train_binary_classifier(
     dataset_repo: str = "geekyrakshit/prompt-injection-dataset",
     model_name: str = "distilbert/distilbert-base-uncased",
     prompt_column_name: str = "prompt",
+    id2label: dict[int, str] = {0: "SAFE", 1: "INJECTION"},
+    label2id: dict[str, int] = {"SAFE": 0, "INJECTION": 1},
     learning_rate: float = 1e-5,
     batch_size: int = 16,
     num_epochs: int = 2,
@@ -49,6 +67,44 @@ def train_binary_classifier(
     save_steps: int = 1000,
     streamlit_mode: bool = False,
 ):
+    """
+    Trains a binary classifier using a specified dataset and model architecture.
+
+    This function sets up and trains a binary sequence classification model using
+    the Hugging Face Transformers library. It integrates with Weights & Biases for
+    experiment tracking and optionally displays a progress bar in a Streamlit app.
+
+    Args:
+        project_name (str): The name of the Weights & Biases project.
+        entity_name (str): The Weights & Biases entity (user or team).
+        run_name (str): The name of the Weights & Biases run.
+        dataset_repo (str, optional): The Hugging Face dataset repository to load.
+            Defaults to "geekyrakshit/prompt-injection-dataset".
+        model_name (str, optional): The pre-trained model to use. Defaults to
+            "distilbert/distilbert-base-uncased".
+        prompt_column_name (str, optional): The column name in the dataset containing
+            the text prompts. Defaults to "prompt".
+        id2label (dict[int, str], optional): Mapping from label IDs to label names.
+            Defaults to {0: "SAFE", 1: "INJECTION"}.
+        label2id (dict[str, int], optional): Mapping from label names to label IDs.
+            Defaults to {"SAFE": 0, "INJECTION": 1}.
+        learning_rate (float, optional): The learning rate for training. Defaults to 1e-5.
+        batch_size (int, optional): The batch size for training and evaluation.
+            Defaults to 16.
+        num_epochs (int, optional): The number of training epochs. Defaults to 2.
+        weight_decay (float, optional): The weight decay for the optimizer. Defaults to 0.01.
+        save_steps (int, optional): The number of steps between model checkpoints.
+            Defaults to 1000.
+        streamlit_mode (bool, optional): If True, integrates with Streamlit to display
+            a progress bar. Defaults to False.
+
+    Returns:
+        dict: The output of the training process, including metrics and model state.
+
+    Raises:
+        Exception: If an error occurs during training, the exception is raised after
+            ensuring Weights & Biases run is finished.
+    """
     wandb.init(project=project_name, entity=entity_name, name=run_name)
     if streamlit_mode:
         st.markdown(
@@ -69,9 +125,6 @@ def train_binary_classifier(
         predictions = np.argmax(predictions, axis=1)
         return accuracy.compute(predictions=predictions, references=labels)
 
-    id2label = {0: "SAFE", 1: "INJECTION"}
-    label2id = {"SAFE": 0, "INJECTION": 1}
-
     model = AutoModelForSequenceClassification.from_pretrained(
         model_name,
         num_labels=2,

mkdocs.yml

@@ -68,6 +68,7 @@ nav:
   - LLM: 'llm.md'
   - Metrics: 'metrics.md'
   - RegexModel: 'regex_model.md'
+  - Train Classifier: 'train_classifier.md'
   - Utils: 'utils.md'
 
 repo_url: https://github.com/soumik12345/guardrails-genie