promptforest 0.1.1__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.1.1.dist-info → promptforest-0.5.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: promptforest
-Version: 0.1.1
+Version: 0.5.0
 Summary: Ensemble Prompt Injection Detection
 Requires-Python: >=3.8
 Description-Content-Type: text/markdown
@@ -36,25 +36,43 @@ This discrepancy score enables downstream workflows such as:
 - 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 (Built with Llama) |
-| **ProtectAI** | DebertaV3 Prompt Injection Finetune       |
-| **Deepset**   | DebertaV3-base Injection Finetune         |
-| **Katanemo**  | Arch-Guard                                |
-| **Appleroll** | PromptForest-XGBoost                      |
-
-## Performance
-**Request Latency** \
-Best Case: 50ms \
-Worst Case: 200ms
-
-**Accuracy** \
-Preliminary results indicate ensemble performance is at least as good as any individual model. Extensive benchmarking is ongoing.
-
+| **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.

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.1.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.1.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.1.dist-info/licenses/LICENSE.txt,sha256=GgVl4CdplCpCEssTcrmIRbz52zQc0fdcSETZp34uBF4,11349
-promptforest-0.1.1.dist-info/licenses/NOTICE.md,sha256=XGjuV5VAWBinW6Jzu7-9h0Ph3xwCNzcJdbMH_EgU_g4,356
-promptforest-0.1.1.dist-info/METADATA,sha256=o1T79TkOnH3uMEWzI31xwmyP-QvFKH2JMHBLFv-WGVI,3700
-promptforest-0.1.1.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-promptforest-0.1.1.dist-info/entry_points.txt,sha256=sVcjABvpA7P2fXca2KMZSYf0PNfDgLt1NHlYFMPO_eE,55
-promptforest-0.1.1.dist-info/top_level.txt,sha256=NxasbbadJaf8w9zaRXo5KOdBqNA1oDe-2X7e6zdz3k0,13
-promptforest-0.1.1.dist-info/RECORD,,