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

promptforest/config.py

@@ -13,10 +13,9 @@ MODELS_DIR = USER_DATA_DIR / "models"
 DEFAULT_CONFIG = {
     "models": [
         {"name": "llama_guard", "path": "llama_guard", "type": "hf", "enabled": True},
-        {"name": "protectai", "path": "protectai_deberta", "type": "hf", "enabled": True},
-        {"name": "deepset", "path": "deepset_deberta", "type": "hf", "enabled": True},
-        {"name": "katanemo", "path": "katanemo_arch", "type": "hf", "enabled": True},
-        {"name": "xgboost", "type": "xgboost", "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}
     ],
     "settings": {
         "device": "auto",  # Options: auto, cuda, mps, cpu

promptforest/download.py

@@ -13,9 +13,8 @@ from .llama_guard_86m_downloader import download_llama_guard
 
 # Configuration
 MODELS = {
-    "protectai_deberta": "protectai/deberta-v3-base-prompt-injection",
-    "deepset_deberta": "deepset/deberta-v3-base-injection",
-    "katanemo_arch": "katanemo/Arch-Guard"
+    "protectai": "protectai/deberta-v3-base-prompt-injection-v2",
+    "vijil_dome": "vijil/vijil_dome_prompt_injection_detection"
 }
 
 EMBEDDING_MODEL_NAME = 'all-MiniLM-L6-v2'
@@ -31,15 +30,20 @@ def _download_hf_model(name, model_id):
     try:
         if save_path.exists():
             return
+        
+        # Special handling for Vijil (ModernBERT tokenizer issue)
+        tokenizer_id = model_id
+        if "vijil" in name or "vijil" in model_id:
+            tokenizer_id = "answerdotai/ModernBERT-base"
 
-        tokenizer = AutoTokenizer.from_pretrained(model_id)
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_id)
         model = AutoModelForSequenceClassification.from_pretrained(model_id)
         
         tokenizer.save_pretrained(save_path)
         model.save_pretrained(save_path)
         
     except Exception as e:
-        print("Failed to download a model!")
+        print(f"Failed to download {model_id}: {e}")
 
 def _download_sentence_transformer():
     """Download and save the SentenceTransformer model."""

promptforest/lib.py

@@ -68,7 +68,8 @@ class HFModel(ModelInference):
         id2label = self.model.config.id2label
         found = False
         for idx, label in id2label.items():
-            if any(kw in label.lower() for kw in MALICIOUS_KEYWORDS):
+            # Check if label is string before calling lower()
+            if isinstance(label, str) and any(kw in label.lower() for kw in MALICIOUS_KEYWORDS):
                 self.malicious_idx = idx
                 found = True
                 break
@@ -99,9 +100,11 @@ class HFModel(ModelInference):
 
 
 class XGBoostModel(ModelInference):
-    def __init__(self, settings):
+    def __init__(self, settings, config=None):
         self.name = "xgboost_custom"
         self.settings = settings
+        self.config = config or {}
+        self.threshold = self.config.get('threshold', 0.5)
         self.model = None
         self.embedder = None
         
@@ -140,7 +143,18 @@ class XGBoostModel(ModelInference):
         try:
             emb = self.embedder.encode([prompt])
             prob = self.model.predict_proba(emb)[0][1]
-            return float(prob)
+            prob = float(prob)
+
+            # Apply threshold adjustment if a custom threshold is set
+            if self.threshold != 0.5:
+                if prob < self.threshold:
+                    # Map [0, threshold) -> [0, 0.5)
+                    prob = 0.5 * (prob / self.threshold)
+                else:
+                    # Map [threshold, 1.0] -> [0.5, 1.0]
+                    prob = 0.5 + 0.5 * (prob - self.threshold) / (1.0 - self.threshold)
+                    
+            return prob
         except Exception:
             return 0.0
 
@@ -151,14 +165,14 @@ class EnsembleGuard:
         Initialize the EnsembleGuard.
         :param config: Dictionary containing configuration. If None, loads default/user config.
         """
-        # Check if models need to be downloaded
-        self._ensure_models_available()
-        
         if config is None:
             self.config = load_config()
         else:
             self.config = config
             
+        # Check if models need to be downloaded
+        self._ensure_models_available()
+            
         self.models = []
         self._init_models()
         self.device_used = get_device(self.config['settings'].get('device', 'auto'))
@@ -167,14 +181,27 @@ class EnsembleGuard:
         """Check if models are available, download if needed."""
         from .config import MODELS_DIR
         
-        # Check if models directory exists and has content
-        if MODELS_DIR.exists() and any(MODELS_DIR.iterdir()):
-            return
-        
-        # Models not found, download them
-        print("Models not found. Downloading...")
-        from .download import download_all
-        download_all()
+        missing = False
+        if not MODELS_DIR.exists():
+            missing = True
+        else:
+            # Check for each enabled HF model
+            for model_cfg in self.config.get('models', []):
+                if model_cfg.get('type') == 'hf' and model_cfg.get('enabled', True):
+                    path = MODELS_DIR / model_cfg['path']
+                    if not path.exists():
+                        missing = True
+                        break
+            
+            # Check for SentenceTransformer (needed for XGBoost)
+            st_path = MODELS_DIR / 'sentence_transformer'
+            if not st_path.exists():
+                missing = True
+
+        if missing:
+            print("Some models not found. Downloading required models...")
+            from .download import download_all
+            download_all()
 
     def _init_models(self):
         settings = self.config.get('settings', {})
@@ -189,7 +216,7 @@ class EnsembleGuard:
             if model_type == 'hf':
                 self.models.append(HFModel(model_cfg['name'], model_cfg['path'], settings))
             elif model_type == 'xgboost':
-                self.models.append(XGBoostModel(settings))
+                self.models.append(XGBoostModel(settings, model_cfg))
             else:
                 print(f"Unknown model type: {model_type}")
 

promptforest/llama_guard_86m_downloader.py

@@ -1,67 +1,58 @@
 """
-Script to download Llama Guard 2 86M from a custom GitHub repository.
-Handles split safetensor files and combines them locally.
+Script to download Llama Guard 2 86M from custom GitHub releases.
+Downloads files in parallel for speed.
 """
 
 import os
-import shutil
-import subprocess
-import tempfile
+import requests
 from pathlib import Path
+from concurrent.futures import ThreadPoolExecutor
 from .config import MODELS_DIR
 
-LLAMA_GUARD_REPO = "https://github.com/appleroll-research/promptforest-model-ensemble.git"
+BASE_URL = "https://github.com/appleroll-research/promptforest-model-ensemble/releases/download/v0.5.0-alpha"
+FILES_TO_DOWNLOAD = [
+    "config.json",
+    "model.safetensors",
+    "special_tokens_map.json",
+    "tokenizer.json",
+    "tokenizer_config.json"
+]
 
-def _download_llama_guard():
-    """Download Llama Guard from custom repository and combine split files."""
-    save_path = MODELS_DIR / "llama_guard"
-    
+def _download_file(url, save_path):
+    """Download a single file."""
     if save_path.exists():
         return
-    
+        
     try:
-        # Use temporary directory for cloning
-        with tempfile.TemporaryDirectory() as temp_dir:
-            temp_path = Path(temp_dir)
-            
-            # Clone repository silently
-            subprocess.run(
-                ["git", "clone", "--depth", "1", LLAMA_GUARD_REPO, str(temp_path)],
-                stdout=subprocess.DEVNULL,
-                stderr=subprocess.DEVNULL,
-                check=True
-            )
-            
-            # Get the llama_guard folder from the cloned repo
-            source_dir = temp_path / "llama_guard"
-            if not source_dir.exists():
-                raise FileNotFoundError(f"llama_guard folder not found in repository")
-            
-            # Copy to models directory
-            save_path.parent.mkdir(parents=True, exist_ok=True)
-            shutil.copytree(source_dir, save_path)
-            
-            # Combine split safetensor files
-            model_file = save_path / "model.safetensors"
-            if not model_file.exists():
-                # Find and combine c_* files
-                split_files = sorted(save_path.glob("c_*"))
-                if split_files:
-                    with open(model_file, 'wb') as outfile:
-                        for split_file in split_files:
-                            with open(split_file, 'rb') as infile:
-                                outfile.write(infile.read())
-                    
-                    # Delete split files
-                    for split_file in split_files:
-                        split_file.unlink()
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
         
+        with open(save_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
     except Exception as e:
-        # Clean up on failure
+        print(f"Failed to download {url}: {e}")
+        # Clean up partial file
         if save_path.exists():
-            shutil.rmtree(save_path)
-        raise Exception(f"Failed to download Llama Guard: {e}")
+            os.remove(save_path)
 
 def download_llama_guard():
-    """Public interface for downloading Llama Guard."""
-    _download_llama_guard()
+    """Download Llama Guard files in parallel."""
+    save_dir = MODELS_DIR / "llama_guard"
+    
+    # Check if all files exist
+    if save_dir.exists() and all((save_dir / f).exists() for f in FILES_TO_DOWNLOAD):
+        return
+
+    save_dir.mkdir(parents=True, exist_ok=True)
+    
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        futures = []
+        for filename in FILES_TO_DOWNLOAD:
+            url = f"{BASE_URL}/{filename}"
+            save_path = save_dir / filename
+            futures.append(executor.submit(_download_file, url, save_path))
+            
+        for future in futures:
+            future.result()
+            

promptforest-0.5.0.dist-info/METADATA

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: promptforest
+Version: 0.5.0
+Summary: Ensemble Prompt Injection Detection
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE.txt
+License-File: NOTICE.md
+Requires-Dist: numpy
+Requires-Dist: pandas
+Requires-Dist: torch
+Requires-Dist: transformers
+Requires-Dist: sentence-transformers
+Requires-Dist: xgboost
+Requires-Dist: scikit-learn
+Requires-Dist: pyyaml
+Requires-Dist: joblib
+Requires-Dist: protobuf
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: license-file
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# PromptForest - Fast and Reliable Injection Detector Ensemble
+
+PromptForest is a prompt injection detector ensemble focused on real-world latency and reliability.
+
+We rely on an ensemble of small, accurate prompt detection models using a voting system to generate accurate detections. 
+
+By comparing predictions across multiple models, the system can flag prompts where models disagree, helping to reduce the risk of false negatives.
+
+This discrepancy score enables downstream workflows such as:
+- Human-in-the-loop review for high-risk or ambiguous prompts  
+- Adaptive throttling or alerting in production systems  
+- Continuous monitoring and model improvement  
+
+## Statistics
+**E2E Request Latency** \
+Average Case: 100ms \
+Worst Case: 200ms
+
+PromptForest was evaluated against the SOTA model Qualifire Sentinel model (v2).
+
+| 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 |
+
+
+### Key Insights
+
+- Calibrated uncertainty: PromptForest is less confident on wrong predictions than Sentinel, 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
+
+| 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.
+
+```bash
+pip install promptforest
+promptforest serve --port 8000 
+```
+
+## Disclaimer & Limitations
+
+PromptForest uses a combination of open-source and third-party machine learning models, including models and weights released by other organizations under their respective licenses (e.g. Meta LLaMA Prompt Guard and other public prompt-injection detectors).
+All third-party components remain the property of their original authors and are used in accordance with their licenses.
+
+PromptForest is not a standalone security solution and should not be relied upon as the sole defense mechanism for protecting production systems. Prompt injection detection is an inherently adversarial and evolving problem, and no automated system can guarantee complete protection.
+
+This project has not yet been extensively validated against real-world, large-scale, or targeted prompt-injection attacks. Results may vary depending on deployment context, model configuration, and threat model.
+
+PromptForest is intended to be used as one layer in a defense-in-depth strategy, alongside input validation, output filtering, access control, sandboxing, monitoring, and human oversight.
+
+## License
+This project is licensed under Apache 2.0. Third-party models and weights are redistributed under their original licenses (see THIRD_PARTY_LICENSES folder for details). Users must comply with these licenses.

promptforest-0.5.0.dist-info/RECORD

@@ -0,0 +1,15 @@
+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,,

{promptforest-0.1.0.dist-info → promptforest-0.5.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.10.1)
+Generator: setuptools (80.10.2)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

promptforest-0.1.0.dist-info/METADATA

@@ -1,21 +0,0 @@
-Metadata-Version: 2.4
-Name: promptforest
-Version: 0.1.0
-Summary: Ensemble Prompt Injection Detection
-Requires-Python: >=3.8
-License-File: LICENSE.txt
-License-File: NOTICE.md
-Requires-Dist: numpy
-Requires-Dist: pandas
-Requires-Dist: torch
-Requires-Dist: transformers
-Requires-Dist: sentence-transformers
-Requires-Dist: xgboost
-Requires-Dist: scikit-learn
-Requires-Dist: pyyaml
-Requires-Dist: joblib
-Requires-Dist: protobuf
-Dynamic: license-file
-Dynamic: requires-dist
-Dynamic: requires-python
-Dynamic: summary

promptforest-0.1.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=bOFHlK0kI7c3fzccZrcjccNUfZJPzvLtKEAZ_loLvzE,3366
-promptforest/download.py,sha256=3Ss1BX6kQatfhif1cbErUekPlSA2RCqtiatUzGi72zo,2454
-promptforest/lib.py,sha256=LT8A1_veV9tB2DyrZ0JEOBW4EWEs9El5xOxF0zNHOAc,8042
-promptforest/llama_guard_86m_downloader.py,sha256=ibFeeuDgMBVe-8aD0zl23xJKOPdKyw-4Bsf0iZJih4s,2412
-promptforest/server.py,sha256=uF4Yj7yR_2vEx_7nQabGHGGw-6GWnT0iBZx3UPQK634,2905
-promptforest/xgboost/xgb_model.pkl,sha256=97Y_Dfu8PwubkplRXJdNEuAj9te1v-nEJlXfPpEZWdM,748772
-promptforest-0.1.0.dist-info/licenses/LICENSE.txt,sha256=GgVl4CdplCpCEssTcrmIRbz52zQc0fdcSETZp34uBF4,11349
-promptforest-0.1.0.dist-info/licenses/NOTICE.md,sha256=XGjuV5VAWBinW6Jzu7-9h0Ph3xwCNzcJdbMH_EgU_g4,356
-promptforest-0.1.0.dist-info/METADATA,sha256=OYvSPhnatbf97rur1W3zaY4FE0MFRE67j8QmC8hpz_M,509
-promptforest-0.1.0.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-promptforest-0.1.0.dist-info/entry_points.txt,sha256=sVcjABvpA7P2fXca2KMZSYf0PNfDgLt1NHlYFMPO_eE,55
-promptforest-0.1.0.dist-info/top_level.txt,sha256=NxasbbadJaf8w9zaRXo5KOdBqNA1oDe-2X7e6zdz3k0,13
-promptforest-0.1.0.dist-info/RECORD,,