Add training dataset switcher

src/ai_security/chatter_demo.py

@@ -3,10 +3,17 @@ import gradio as gr
 
 from .generative_chatter_detector import GenerativeChatterDetector
 from .rules_chatter_detector import  simple_normalized_blacbriar_chatter_detector
-from .discriminative_chatter_detector import DiscriminativeChatterDetector
+from .discriminative_chatter_detector import DatasetName, DiscriminativeChatterDetector
+
+
+
+def update_dataset(selected_dataset):
+    """Reinitialize discriminative detector with new dataset"""
+    global discriminative_detector
+    discriminative_detector = DiscriminativeChatterDetector(dataset_name=selected_dataset)
+    return f"✓ Classic ML model now trained on: {selected_dataset}"
 
 
-discriminative_detector = DiscriminativeChatterDetector(scope='blackbriar-only')
 generative_detector = GenerativeChatterDetector(scope='blackbriar-only')
 
 
@@ -57,14 +64,26 @@ with gr.Blocks(theme=hwr_theme) as demo:
     """)
 
     with gr.Row():
-        input_text = gr.Textbox(
-            label="Enter text to analyze",
-            placeholder="Type your message here...",
-            lines=3
+        with gr.Column(scale=2):
+            input_text = gr.Textbox(
+                label="Enter text to analyze",
+                placeholder="Type your message here...",
+                lines=3
+            )
+        with gr.Column(scale=1):
+            dataset_name = gr.Dropdown(
+            choices=[a_name.value for a_name in DatasetName],
+            value='blackbriar',  # Set default value
+            label="(Optional) Change training Dataset for Classic ML"
         )
+        dataset_status = gr.Markdown("Currently using: blackbriar")
+
 
     gr.Markdown("### Compare All Three Models")
 
+    # with gr.Row():
+
+
     with gr.Row():
         with gr.Column():
             gr.Markdown("#### Model A")
@@ -81,7 +100,6 @@ with gr.Blocks(theme=hwr_theme) as demo:
             output_c = gr.Markdown()
             btn_c = gr.Button("Analyze with Model C", variant="primary")
 
-    # Example inputs
     gr.Examples(
         examples=[
             [" Bourne's just the tip of the iceberg. Have you heard of an 'Operation Blackbriar'?"],
@@ -93,8 +111,13 @@ with gr.Blocks(theme=hwr_theme) as demo:
     )
     clear_btn = gr.Button("Let's try again.")
 
-    # Connect buttons to functions
     btn_a.click(fn=detect_chatter_a, inputs=input_text, outputs=output_a)
     btn_b.click(fn=detect_chatter_b, inputs=input_text, outputs=output_b)
     btn_c.click(fn=detect_chatter_c, inputs=input_text, outputs=output_c)
-    clear_btn.click(fn=clear_all_outputs, inputs=None, outputs=[input_text, output_a, output_b, output_c])
+    clear_btn.click(fn=clear_all_outputs, inputs=None, outputs=[input_text, output_a, output_b, output_c])
+
+    dataset_name.change(
+        fn=update_dataset,
+        inputs=dataset_name,
+        outputs=dataset_status
+    )

src/ai_security/discriminative_chatter_detector.py

@@ -1,19 +1,27 @@
 """
-Chatter detection using bag-of-words and multinomial bayes
+Chatter detection using discriminative ML
 """
-import pandas as pd
+from enum import Enum
 from pathlib import Path
+
+import pandas as pd
 from sklearn.feature_extraction.text import CountVectorizer
 from sklearn.linear_model import LogisticRegression
 from sklearn.pipeline import Pipeline
 
 
-DATASET_NAME_FILENAME_MAPPING = {
-    'blackbriar': 'chatter-detection-dataset - blackbriar-chatter-detection-dataset.csv',
-    'extended': 'chatter-detection-dataset - extended-chatter-detection-dataset.csv',
-    'baby-blackbriar': 'chatter-detection-dataset - baby-chatter-detection-dataset.csv',
-    'toddler-blackbriar': 'chatter-detection-dataset - toddler-chatter-detection-dataset.csv'
+class DatasetName(Enum):
+    baby_blackbriar = 'baby-blackbriar'
+    toddler_blackbriar = 'toddler-blackbriar'
+    blackbriar = 'blackbriar'
+    extended = 'extended'
+
 
+DATASET_NAME_FILENAME_MAPPING = {
+    DatasetName.baby_blackbriar.value : 'chatter-detection-dataset - baby-chatter-detection-dataset.csv',
+    DatasetName.toddler_blackbriar.value: 'chatter-detection-dataset - toddler-chatter-detection-dataset.csv',
+    DatasetName.blackbriar.value: 'chatter-detection-dataset - blackbriar-chatter-detection-dataset.csv',
+    DatasetName.extended.value: 'chatter-detection-dataset - extended-chatter-detection-dataset.csv',
 }
 
 

src/ai_security/malware_demo.py

@@ -5,7 +5,7 @@ from .discriminative_chatter_detector import DiscriminativeChatterDetector
 from .generative_malware_detector import GenerativeMalwareDetector
 
 
-detector_a = DiscriminativeChatterDetector(scope='blackbriar-only')
+detector_a = DiscriminativeChatterDetector(dataset_name='blackbriar')
 detector_c = GenerativeMalwareDetector()