promptforest 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

promptforest/__init__.py

@@ -1,3 +1,3 @@
 __version__ = "0.1.0"
 
-from .lib import EnsembleGuard
+from .lib import PFEnsemble

promptforest/cli.py

@@ -33,11 +33,11 @@ def main():
         run_server(port=args.port, config=cfg)
 
     elif args.command == "check":
-        from .lib import EnsembleGuard
+        from .lib import PFEnsemble
         cfg = get_user_config(args.config)
         try:
             print(f"Loading PromptForest...")
-            guard = EnsembleGuard(config=cfg)
+            guard = PFEnsemble(config=cfg)
             print(f"Device: {guard.device_used}")
             result = guard.check_prompt(args.prompt)
             print(json.dumps(result, indent=2))

promptforest/config.py

@@ -12,10 +12,9 @@ MODELS_DIR = USER_DATA_DIR / "models"
 
 DEFAULT_CONFIG = {
     "models": [
-        {"name": "llama_guard", "path": "llama_guard", "type": "hf", "enabled": True},
-        {"name": "protectai", "path": "protectai", "type": "hf", "enabled": True},
-        {"name": "vijil", "path": "vijil_dome", "type": "hf", "enabled": True},
-        {"name": "xgboost", "type": "xgboost", "enabled": True, "threshold": 0.10}
+        {"name": "llama_guard", "path": "llama_guard", "type": "hf", "enabled": True, "accuracy_weight": 1.0},
+        {"name": "vijil", "path": "vijil_dome", "type": "hf", "enabled": True, "accuracy_weight": 1.0},
+        {"name": "xgboost", "type": "xgboost", "enabled": True, "threshold": 0.10, "accuracy_weight": 1.0}
     ],
     "settings": {
         "device": "auto",  # Options: auto, cuda, mps, cpu
@@ -30,7 +29,7 @@ def load_config(config_path=None):
     """
     Load configuration from a YAML file, merging with defaults.
     """
-    # Start with a deep copy of the default config structure
+    # Deep copy of the default config structure
     config = {
         "models": [m.copy() for m in DEFAULT_CONFIG["models"]],
         "settings": DEFAULT_CONFIG["settings"].copy(),
@@ -44,19 +43,17 @@ def load_config(config_path=None):
                 with open(path, 'r') as f:
                     user_config = yaml.safe_load(f)
                     if user_config:
-                        # 1. Merge Settings
+                        # Merge config
+                        # @todo: is there a smarter way to merge this?
                         if "settings" in user_config:
                             config["settings"].update(user_config["settings"])
                             
-                        # 2. Merge Logging
                         if "logging" in user_config:
                             config["logging"].update(user_config["logging"])
                             
-                        # 3. Merge Models (Smart Merge)
                         if "models" in user_config:
                             user_models = user_config["models"]
                             if isinstance(user_models, list):
-                                # Convert current models to dict for easy lookup by name
                                 existing_model_map = {m["name"]: m for m in config["models"]}
                                 
                                 for u_model in user_models:

promptforest/download.py

@@ -11,9 +11,9 @@ from sentence_transformers import SentenceTransformer
 from .config import MODELS_DIR
 from .llama_guard_86m_downloader import download_llama_guard
 
-# Configuration
 MODELS = {
-    "protectai": "protectai/deberta-v3-base-prompt-injection-v2",
+    # We are currently not using ProtectAI as it actively degrades ensemble performance while being 86M params
+    # "protectai": "protectai/deberta-v3-base-prompt-injection-v2",
     "vijil_dome": "vijil/vijil_dome_prompt_injection_detection"
 }
 
@@ -31,8 +31,11 @@ def _download_hf_model(name, model_id):
         if save_path.exists():
             return
         
-        # Special handling for Vijil (ModernBERT tokenizer issue)
+        
         tokenizer_id = model_id
+
+        # Vijil uses ModernBERT tokenizer
+        # @todo this should not be hardcoded
         if "vijil" in name or "vijil" in model_id:
             tokenizer_id = "answerdotai/ModernBERT-base"
 
@@ -47,17 +50,14 @@ def _download_hf_model(name, model_id):
 
 def _download_sentence_transformer():
     """Download and save the SentenceTransformer model."""
-    # print(f"Downloading SentenceTransformer ({EMBEDDING_MODEL_NAME})...")
     save_path = MODELS_DIR / 'sentence_transformer'
     
     try:
         if save_path.exists():
-            #  print(f"  - Model already exists at {save_path}. Skipping.")
              return
 
         model = SentenceTransformer(EMBEDDING_MODEL_NAME)
         model.save(str(save_path))
-        #print(f"  - Saved to {save_path}")
         
     except Exception as e:
         print(f"SentenceTransformer download failed: {e}")
@@ -66,11 +66,11 @@ def download_all():
     print(f"Downloading models to {MODELS_DIR}...")
     _ensure_dir(MODELS_DIR)
     
-    # Download Llama Guard in parallel (slowest download)
+    # Download Llama Guard first as it takes the longest
     llama_guard_thread = threading.Thread(target=download_llama_guard, daemon=False)
     llama_guard_thread.start()
     
-    # Download HF Classification Models
+    # Download each model from Hugging Face
     for name, model_id in MODELS.items():
         _download_hf_model(name, model_id)
         

promptforest/lib.py

@@ -16,7 +16,7 @@ from .config import MODELS_DIR, XGB_MODEL_PATH, load_config
 
 # Suppress Warnings
 transformers_logging.set_verbosity_error()
-os.environ["TOKENIZERS_PARALLELISM"] = "false" # Prevent deadlocks/warnings
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 MALICIOUS_KEYWORDS = ['unsafe', 'malicious', 'injection', 'attack', 'jailbreak']
 
@@ -45,7 +45,7 @@ class HFModel(ModelInference):
 
     def _load(self):
         if not self.path.exists():
-             print(f"[WARN] Model path not found: {self.path}")
+             print(f"Model path not found: {self.path}")
              return
 
         try:
@@ -61,7 +61,7 @@ class HFModel(ModelInference):
             self.model.eval()
             self._determine_label_map()
         except Exception as e:
-            print(f"[ERR] Failed to load {self.name}: {e}")
+            print(f"Error: Failed to load {self.name}: {e}")
             self.model = None
 
     def _determine_label_map(self):
@@ -101,9 +101,9 @@ class HFModel(ModelInference):
 
 class XGBoostModel(ModelInference):
     def __init__(self, settings, config=None):
-        self.name = "xgboost_custom"
-        self.settings = settings
         self.config = config or {}
+        self.name = self.config.get("name", "xgboost")
+        self.settings = settings
         self.threshold = self.config.get('threshold', 0.5)
         self.model = None
         self.embedder = None
@@ -123,7 +123,6 @@ class XGBoostModel(ModelInference):
                 self.embedder = SentenceTransformer(str(ST_PATH))
             else:
                 print("Cannot find local SentenceTransformer model. Downloading...")
-                # Fallback to download default if local cache missing
                 self.embedder = SentenceTransformer('all-MiniLM-L6-v2')
             
             if self.device_name in ['cuda', 'mps']:
@@ -134,7 +133,7 @@ class XGBoostModel(ModelInference):
                    except:
                        pass
         except Exception as e:
-            print(f"[ERR] Failed to load XGBoost: {e}")
+            print(f"Error: Failed to load XGBoost: {e}")
             self.model = None
 
     def predict(self, prompt):
@@ -159,10 +158,10 @@ class XGBoostModel(ModelInference):
             return 0.0
 
 
-class EnsembleGuard:
+class PFEnsemble:
     def __init__(self, config=None):
         """
-        Initialize the EnsembleGuard.
+        Initialize the PFEnsemble.
         :param config: Dictionary containing configuration. If None, loads default/user config.
         """
         if config is None:
@@ -221,6 +220,11 @@ class EnsembleGuard:
                 print(f"Unknown model type: {model_type}")
 
     def check_prompt(self, prompt):
+        """
+        Checks the prompt using the ensemble of models.
+        
+        :param prompt: The prompt string to check.
+        """
         start_time = time.perf_counter()
         results = {}
         
@@ -241,24 +245,43 @@ class EnsembleGuard:
         if not probs:
              return {"error": "No models loaded"}
 
+        # Calculate weighted average
+        weighted_sum = 0.0
+        total_weight = 0.0
+        model_configs = {m['name']: m for m in self.config.get('models', [])}
+        
+        for model_name, prob in results.items():
+            model_cfg = model_configs.get(model_name, {})
+            # Default weight is 1.0 if not specified
+            weight = float(model_cfg.get('accuracy_weight', 1.0))
+            weighted_sum += prob * weight
+            total_weight += weight
+            
+        if total_weight > 0:
+            w_avg_prob = weighted_sum / total_weight
+        else:
+            w_avg_prob = np.mean(probs)
+
         avg_prob = np.mean(probs)
+
         max_prob = np.max(probs)
         
         # Uncertainty
         std_dev = np.std(probs)
         uncertainty_score = min(std_dev * 2, 1.0)
         
-        is_malicious = avg_prob > 0.5
+        is_malicious = w_avg_prob > 0.5
         
         response = {
             "is_malicious": bool(is_malicious),
+            # We use average probability for confidence - better results (2-3x improvement in benchmarks)
             "confidence": float(avg_prob if is_malicious else 1 - avg_prob), 
             "uncertainty": float(uncertainty_score),
             "malicious_score": float(avg_prob),
             "max_risk_score": float(max_prob)
         }
         
-        # Add stats if requested
+        # Add stats if logging is requested
         if self.config.get('logging', {}).get('stats', True):
             response["details"] = results
             response["latency_ms"] = duration_ms

promptforest/llama_guard_86m_downloader.py

@@ -1,6 +1,7 @@
 """
 Script to download Llama Guard 2 86M from custom GitHub releases.
-Downloads files in parallel for speed.
+
+@todo: This should be legal, but as soon as anyone in Meta hints us to remove this, we will comply immediately.
 """
 
 import os
@@ -42,6 +43,7 @@ def download_llama_guard():
     
     # Check if all files exist
     if save_dir.exists() and all((save_dir / f).exists() for f in FILES_TO_DOWNLOAD):
+        # All files already exist, we don't need to download them again
         return
 
     save_dir.mkdir(parents=True, exist_ok=True)

promptforest/server.py

@@ -1,5 +1,5 @@
 """
-Simple HTTP Server for PromptForest.
+Server module for PromptForest. Sets up an HTTP server to handle inference requests
 """
 
 import http.server
@@ -7,12 +7,12 @@ import socketserver
 import json
 import sys
 import time
-from .lib import EnsembleGuard
+from .lib import PFEnsemble
 
 PORT = 8000
 ensemble = None
 
-class GuardRequestHandler(http.server.BaseHTTPRequestHandler):
+class PFRequestHandler(http.server.BaseHTTPRequestHandler):
     def do_POST(self):
         """Handle POST requests for inference."""
         if self.path == '/analyze':
@@ -65,7 +65,7 @@ def run_server(port=8000, config=None):
     global ensemble
     print(f"Initializing PromptForest...")
     try:
-        ensemble = EnsembleGuard(config=config)
+        ensemble = PFEnsemble(config=config)
         print(f"Device: {ensemble.device_used}")
         print("Warming up...")
         ensemble.check_prompt("warmup")
@@ -77,7 +77,7 @@ def run_server(port=8000, config=None):
     print(f"\nListening on http://localhost:{port}")
     socketserver.TCPServer.allow_reuse_address = True
     
-    with ThreadedHTTPServer(("", port), GuardRequestHandler) as httpd:
+    with ThreadedHTTPServer(("", port), PFRequestHandler) as httpd:
         try:
             httpd.serve_forever()
         except KeyboardInterrupt:

{promptforest-0.5.0.dist-info → promptforest-0.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: promptforest
-Version: 0.5.0
+Version: 0.7.0
 Summary: Ensemble Prompt Injection Detection
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
@@ -24,6 +24,10 @@ Dynamic: requires-python
 Dynamic: summary
 
 # PromptForest - Fast and Reliable Injection Detector Ensemble
+![PyPI Downloads](https://img.shields.io/pypi/dm/promptforest)
+![Apache License](https://img.shields.io/badge/license-Apache%20License%202.0-blue)
+
+**📖 TRY IT OUT ON A NOTEBOOK [HERE](https://colab.research.google.com/drive/1EW49Qx1ZlaAYchqplDIVk2FJVzCqOs6B?usp=sharing)!**
 
 PromptForest is a prompt injection detector ensemble focused on real-world latency and reliability.
 
@@ -36,53 +40,52 @@ This discrepancy score enables downstream workflows such as:
 - Adaptive throttling or alerting in production systems  
 - Continuous monitoring and model improvement  
 
+## Quick Start
+To use PromptForest, simply install the pip package and serve it at a port of your choice. It should automatically start downloading the default model ensemble.
+
+Gated models are downloaded through our own [ensemble github respository](https://github.com/appleroll-research/promptforest-model-ensemble) and are released in accordance to their terms and conditions.
+
+```bash
+pip install promptforest
+promptforest serve --port 8000 
+```
+
 ## Statistics
 **E2E Request Latency** \
 Average Case: 100ms \
 Worst Case: 200ms
 
-PromptForest was evaluated against the SOTA model Qualifire Sentinel model (v2).
+PromptForest was evaluated against the models from Deepset, ProtectAI, Meta and Vijil, with Promptforest and the SOTA model Qualifire's Sentinel V2 performing the best in terms of reliability.
 
 | Metric                           | PromptForest | Sentinel v2 |
 | -------------------------------- | ------------ | ----------- |
-| Accuracy                         | 0.802        | 0.982       |
-| Avg Confidence on Wrong Answers  | 0.643        | 0.858       |
-| Expected Calibration Error (ECE) | 0.060        | 0.202       |
-| Approximate Model Size           | ~250M params  | 600M params |
+| Accuracy                         | 0.901        | 0.973       |
+| Avg Confidence on Wrong Answers  | 0.642        | 0.76       |
+| Expected Calibration Error (ECE) | 0.070        | 0.096       |
+| Total Model Size           | ~237M params  | ~600M params |
 
 
 ### Key Insights
 
-- Calibrated uncertainty: PromptForest is less confident on wrong predictions than Sentinel, resulting in a much lower ECE.
+- Calibrated uncertainty: PromptForest is less confident on wrong predictions than compared models, resulting in a much lower ECE.
 
 - Parameter efficiency: Achieves competitive reliability with <50% of the parameters.
 
 - Interpretability: Confidence scores can be used to flag uncertain predictions for human review.
 
-Interpretation:
-While Sentinel has higher raw accuracy, PromptForest provides better-calibrated confidence. For systems where overconfidence on wrong answers is risky, PromptForest can reduce the chance of critical errors despite being smaller and faster. 
-
 Using Sentinel v2 as a baseline, and given that other models perform worse than Sentinel in published benchmarks, PromptForest is expected to offer more reliable and calibrated predictions than most alternatives.
 
 
-## Supported Models
+## Models
 
 | Provider      | Model Name                 |
 | ------------- | ----------------------------------------- |
 | **Meta**      | [Llama Prompt Guard 86M](https://huggingface.co/meta-llama/Prompt-Guard-86M) (Built with Llama) |
-| **ProtectAI** | [DebertaV3 Prompt Injection Finetune](https://huggingface.co/protectai/deberta-v3-base-prompt-injection-v2)       |
 | **Vijil**     | [Vijil Dome Prompt Injection Detection](https://huggingface.co/vijil/vijil_dome_prompt_injection_detection)     |
-| **Appleroll** | [PromptForest-XGB](appleroll/promptforest-xgb)                      |
-
-## Quick Start
-To use PromptForest, simply install the pip package and serve it at a port of your choice. It should automatically start downloading the default model ensemble.
-
-Gated models are downloaded through our own [ensemble github respository](https://github.com/appleroll-research/promptforest-model-ensemble) and are released in accordance to their terms and conditions.
+| **Appleroll** | [PromptForest-XGB](https://huggingface.co/appleroll/promptforest-xgb)                      |
 
-```bash
-pip install promptforest
-promptforest serve --port 8000 
-```
+## Current Goals
+This project is actively being updated. Our current focus is on implementing weights on individual models to improve accuracy, as well as retraining the XGBoost model with an updated corpus.
 
 ## Disclaimer & Limitations
 

promptforest-0.7.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+promptforest/__init__.py,sha256=wbIBy-XFARDVm5TmFud-IBumHPZ5ps8Phsjz0tYTUgU,51
+promptforest/cli.py,sha256=T4sLBrNp09mHePOBHSB4QobyCDZdbLcBEm3_saoRMQU,1850
+promptforest/config.py,sha256=MyeH2qeHpE3mVmqaSrjC8VP8P0Q0ZsftdN3TMfQGUP4,3138
+promptforest/download.py,sha256=lKje_L2-CU2e56U932l9Q3ueLDt-Mcq8SLEMNOk_lBA,2572
+promptforest/lib.py,sha256=Sp8DdvTooKJVGqG2MtFaFhoGbEgACHxyVxRCjsALlXs,10053
+promptforest/llama_guard_86m_downloader.py,sha256=MS9YG6MepU0ToskZ9f2iwESg2EnXKGyFzWXwR2s2Xac,1866
+promptforest/server.py,sha256=NY6mn9l4-PpTvkQx6zlQpoE6vDTv_Cm1QvwE3SfWi6g,2940
+promptforest/xgboost/xgb_model.pkl,sha256=kSG2r-6TGfhNJfzwklLQOSgG2z610Z5BXxtgQdXE8Vk,2116991
+promptforest-0.7.0.dist-info/licenses/LICENSE.txt,sha256=GgVl4CdplCpCEssTcrmIRbz52zQc0fdcSETZp34uBF4,11349
+promptforest-0.7.0.dist-info/licenses/NOTICE.md,sha256=XGjuV5VAWBinW6Jzu7-9h0Ph3xwCNzcJdbMH_EgU_g4,356
+promptforest-0.7.0.dist-info/METADATA,sha256=99z2XiN2SFQVons3IHfZW6Iai7ndHNt6XkEk8LWTmb8,5222
+promptforest-0.7.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+promptforest-0.7.0.dist-info/entry_points.txt,sha256=sVcjABvpA7P2fXca2KMZSYf0PNfDgLt1NHlYFMPO_eE,55
+promptforest-0.7.0.dist-info/top_level.txt,sha256=NxasbbadJaf8w9zaRXo5KOdBqNA1oDe-2X7e6zdz3k0,13
+promptforest-0.7.0.dist-info/RECORD,,

promptforest-0.5.0.dist-info/RECORD

@@ -1,15 +0,0 @@
-promptforest/__init__.py,sha256=cE1cQyRL4vUzseCwLYbI5wrZuZ-NRMVXIjAgwTLwIEs,54
-promptforest/cli.py,sha256=LKsnbEQNQ9pP_Ww24Ql2Tb_uomO-StqHnk-IHONSKTM,1856
-promptforest/config.py,sha256=c_7GX7nh_1Aa-QU7SOZlthPNGXSoh2KvYOk7txJeQh4,3284
-promptforest/download.py,sha256=6TQvo2qd3tUUxJU6MMsFMgOciHP5HNDNEo3UTOeYI34,2637
-promptforest/lib.py,sha256=WEuEhNNlRQAerLyEIbTHdi15qdXUMuiQOhfsvaftj4M,9254
-promptforest/llama_guard_86m_downloader.py,sha256=0B2ttwLWHki0yLEoJG3BwyFE1oqJFY0M2mLEtmMWmPk,1720
-promptforest/server.py,sha256=uF4Yj7yR_2vEx_7nQabGHGGw-6GWnT0iBZx3UPQK634,2905
-promptforest/xgboost/xgb_model.pkl,sha256=kSG2r-6TGfhNJfzwklLQOSgG2z610Z5BXxtgQdXE8Vk,2116991
-promptforest-0.5.0.dist-info/licenses/LICENSE.txt,sha256=GgVl4CdplCpCEssTcrmIRbz52zQc0fdcSETZp34uBF4,11349
-promptforest-0.5.0.dist-info/licenses/NOTICE.md,sha256=XGjuV5VAWBinW6Jzu7-9h0Ph3xwCNzcJdbMH_EgU_g4,356
-promptforest-0.5.0.dist-info/METADATA,sha256=fEgp4u7q-P74Zo3eF0gnEjVSFMuIc9z9g-1AoAKPAZs,5002
-promptforest-0.5.0.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-promptforest-0.5.0.dist-info/entry_points.txt,sha256=sVcjABvpA7P2fXca2KMZSYf0PNfDgLt1NHlYFMPO_eE,55
-promptforest-0.5.0.dist-info/top_level.txt,sha256=NxasbbadJaf8w9zaRXo5KOdBqNA1oDe-2X7e6zdz3k0,13
-promptforest-0.5.0.dist-info/RECORD,,