Swap around models

app.py

@@ -5,8 +5,8 @@ from ai_security.chatter_demo import demo as chatter_demo
 from ai_security.malware_demo import demo as malware_demo
 
 
-# demo = chatter_demo
-demo = malware_demo
+demo = chatter_demo
+# demo = malware_demo
 
 if __name__ == "__main__":
 

src/ai_security/chatter_demo.py

@@ -6,13 +6,13 @@ from .rules_chatter_detector import  simple_normalized_blacbriar_chatter_detecto
 from .discriminative_chatter_detector import DiscriminativeChatterDetector
 
 
-detector_a = DiscriminativeChatterDetector(scope='blackbriar-only')
-detector_c = GenerativeChatterDetector(scope='blackbriar-only')
+discriminative_detector = DiscriminativeChatterDetector(scope='blackbriar-only')
+generative_detector = GenerativeChatterDetector(scope='blackbriar-only')
 
 
 def detect_chatter_a(text):
     """Model A detection"""
-    result = detector_a.predict(text)
+    result = generative_detector.predict(text)
     return f"**Prediction:** {result['label']}"
 
 
@@ -26,7 +26,7 @@ def detect_chatter_b(text):
 
 def detect_chatter_c(text):
     """Model C detection"""
-    result = detector_c.detect(text)
+    result = discriminative_detector.predict(text)
     return f"**Prediction:** {result['label']}"
 
 

src/ai_security/generative_chatter_detector.py

@@ -26,7 +26,7 @@ class GenerativeChatterDetector:
         self.client = InferenceClient(model='google/gemma-2-2b-it', token=os.environ['HF_INFERENCE_TOKEN'])
 
 
-    def detect(self, text):
+    def predict(self, text):
         try:
             prompt = self.prompt_template.format(text=text)
             print(f"Prompt: {prompt}")