add model toggle: DistilBERT v3 + HRM-Text with ONNX Runtime Web

index.html

@@ -74,6 +74,48 @@
     background: var(--accent);
   }
 
+  .model-toggle {
+    display: flex;
+    gap: 0.5rem;
+    margin-bottom: 1rem;
+  }
+
+  .model-toggle-btn {
+    flex: 1;
+    padding: 0.75rem 0.75rem;
+    background: var(--surface);
+    border: 1px solid var(--border);
+    border-radius: 10px;
+    color: var(--text-dim);
+    font-size: 0.8125rem;
+    cursor: pointer;
+    transition: border-color 0.15s, color 0.15s, background 0.15s;
+    text-align: left;
+    line-height: 1.5;
+  }
+
+  .model-toggle-btn:hover {
+    border-color: color-mix(in srgb, var(--accent) 50%, transparent);
+  }
+
+  .model-toggle-btn.active {
+    border-color: var(--accent);
+    color: var(--text);
+    background: color-mix(in srgb, var(--accent) 8%, var(--surface));
+  }
+
+  .model-toggle-btn .toggle-name {
+    font-weight: 600;
+    font-size: 0.875rem;
+    display: block;
+    margin-bottom: 0.125rem;
+  }
+
+  .model-toggle-btn .toggle-meta {
+    font-size: 0.6875rem;
+    opacity: 0.7;
+  }
+
   .input-area {
     position: relative;
   }
@@ -111,7 +153,7 @@
     font-variant-numeric: tabular-nums;
   }
 
-  button {
+  button.analyze-btn {
     padding: 0.625rem 1.5rem;
     background: var(--accent);
     color: #fff;
@@ -123,8 +165,8 @@
     transition: opacity 0.15s;
   }
 
-  button:hover { opacity: 0.9; }
-  button:disabled { opacity: 0.4; cursor: not-allowed; }
+  button.analyze-btn:hover { opacity: 0.9; }
+  button.analyze-btn:disabled { opacity: 0.4; cursor: not-allowed; }
 
   .result {
     margin-top: 1.5rem;
@@ -287,19 +329,30 @@
 <div class="loading-overlay" id="loading">
   <div class="spinner"></div>
   <div class="loading-text" id="loading-text">Loading model...</div>
-  <div class="loading-detail">~65 MB quantized DistilBERT (one-time download)</div>
+  <div class="loading-detail" id="loading-detail">~65 MB quantized DistilBERT (one-time download)</div>
 </div>
 
 <div class="container">
   <header>
     <h1>Prompt Injection Detector</h1>
-    <p class="subtitle">Detects prompt injection attacks in text using a DistilBERT model fine-tuned on 476K adversarial samples. Runs entirely in your browser.</p>
+    <p class="subtitle">Detects prompt injection attacks in text using ML models running entirely in your browser.</p>
     <div class="model-badge">
       <span class="dot"></span>
       <span id="status-text">Loading...</span>
     </div>
   </header>
 
+  <div class="model-toggle">
+    <button class="model-toggle-btn active" id="toggle-distilbert" data-model="distilbert">
+      <span class="toggle-name">DistilBERT v3</span>
+      <span class="toggle-meta">67M params &middot; F1 0.9961 &middot; mixed data</span>
+    </button>
+    <button class="model-toggle-btn" id="toggle-hrm" data-model="hrm">
+      <span class="toggle-name">HRM-Text</span>
+      <span class="toggle-meta">46.2M params &middot; F1 0.9886 &middot; byte-level</span>
+    </button>
+  </div>
+
   <div class="input-area">
     <textarea
       id="input"
@@ -308,7 +361,7 @@
     ></textarea>
     <div class="controls">
       <span class="char-count" id="char-count">0 chars</span>
-      <button id="analyze" disabled>Analyze</button>
+      <button class="analyze-btn" id="analyze" disabled>Analyze</button>
     </div>
   </div>
 
@@ -336,20 +389,25 @@
   </div>
 
   <footer>
-    <p>Model: <a href="https://huggingface.co/av-codes/prompt-injection-detector-v2-bordair">av-codes/prompt-injection-detector-v2-bordair</a></p>
-    <p>DistilBERT 67M · F1 0.9993 · Trained on <a href="https://huggingface.co/datasets/Bordair/bordair-multimodal">bordair-multimodal</a> (476K samples, v1-v5 adversarial attacks)</p>
-    <p>Powered by <a href="https://huggingface.co/docs/transformers.js">Transformers.js</a> · Inference runs locally in your browser</p>
+    <p><strong>DistilBERT v3:</strong> <a href="https://huggingface.co/av-codes/prompt-injection-detector-v3-mixed">av-codes/prompt-injection-detector-v3-mixed</a> &middot; 67M params &middot; F1 0.9961 &middot; mixed training data (bordair + v1)</p>
+    <p><strong>HRM-Text:</strong> <a href="https://huggingface.co/av-codes/prompt-injection-hrm-text">av-codes/prompt-injection-hrm-text</a> &middot; 46.2M params &middot; F1 0.9886 &middot; from-scratch byte-level &middot; bordair data</p>
+    <p>Powered by <a href="https://huggingface.co/docs/transformers.js">Transformers.js</a> + <a href="https://onnxruntime.ai/">ONNX Runtime Web</a> &middot; Inference runs locally in your browser</p>
   </footer>
 </div>
 
 <script type="module">
   import { pipeline, env } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.0';
+  import * as ort from 'https://cdn.jsdelivr.net/npm/onnxruntime-web@1.21.0/dist/ort.min.mjs';
 
   env.allowLocalModels = false;
 
-  const MODEL_ID = 'av-codes/prompt-injection-detector-v2-bordair';
+  const DISTILBERT_MODEL_ID = 'av-codes/prompt-injection-detector-v3-mixed';
+  const HRM_ONNX_URL = 'https://huggingface.co/av-codes/prompt-injection-hrm-text/resolve/main/onnx/model_fp16.onnx';
+  const HRM_MAX_LEN = 2048;
+
   const loadingEl = document.getElementById('loading');
   const loadingText = document.getElementById('loading-text');
+  const loadingDetail = document.getElementById('loading-detail');
   const statusText = document.getElementById('status-text');
   const inputEl = document.getElementById('input');
   const analyzeBtn = document.getElementById('analyze');
@@ -359,54 +417,177 @@
   const resultConfidence = document.getElementById('result-confidence');
   const confidenceFill = document.getElementById('confidence-fill');
   const latencyEl = document.getElementById('latency');
+  const toggleDistilbert = document.getElementById('toggle-distilbert');
+  const toggleHrm = document.getElementById('toggle-hrm');
+
+  let activeModel = 'distilbert';
+  let distilbertClassifier = null;
+  let hrmSession = null;
+  let isLoading = false;
+
+  function showLoading(modelName, sizeInfo) {
+    loadingText.textContent = `Loading ${modelName}...`;
+    loadingDetail.textContent = sizeInfo;
+    loadingEl.classList.remove('hidden');
+  }
 
-  let classifier = null;
+  function hideLoading() {
+    loadingEl.classList.add('hidden');
+  }
+
+  function updateBadge(text) {
+    statusText.textContent = text;
+  }
 
-  async function loadModel() {
+  async function loadDistilbert() {
+    if (distilbertClassifier) return;
+    showLoading('DistilBERT v3', '~65 MB quantized (one-time download)');
     try {
-      loadingText.textContent = 'Downloading model...';
-      classifier = await pipeline('text-classification', MODEL_ID, {
+      distilbertClassifier = await pipeline('text-classification', DISTILBERT_MODEL_ID, {
         dtype: 'q8',
         device: 'wasm',
         progress_callback: (progress) => {
           if (progress.status === 'progress' && progress.total) {
             const pct = Math.round((progress.loaded / progress.total) * 100);
-            loadingText.textContent = `Downloading... ${pct}%`;
+            loadingText.textContent = `Downloading DistilBERT v3... ${pct}%`;
           } else if (progress.status === 'ready') {
-            loadingText.textContent = 'Model ready';
+            loadingText.textContent = 'DistilBERT v3 ready';
           }
         }
       });
+    } catch (err) {
+      loadingText.textContent = `Error: ${err.message}`;
+      console.error(err);
+      throw err;
+    }
+  }
 
-      loadingEl.classList.add('hidden');
-      statusText.textContent = 'Ready';
-      analyzeBtn.disabled = false;
+  async function loadHrm() {
+    if (hrmSession) return;
+    showLoading('HRM-Text', '~94 MB ONNX (one-time download)');
+    try {
+      loadingText.textContent = 'Downloading HRM-Text...';
+      hrmSession = await ort.InferenceSession.create(HRM_ONNX_URL, {
+        executionProviders: ['wasm'],
+      });
     } catch (err) {
       loadingText.textContent = `Error: ${err.message}`;
       console.error(err);
+      throw err;
+    }
+  }
+
+  function softmax(logits) {
+    const max = Math.max(...logits);
+    const exps = logits.map(x => Math.exp(x - max));
+    const sum = exps.reduce((a, b) => a + b, 0);
+    return exps.map(x => x / sum);
+  }
+
+  function tokenizeBytes(text) {
+    const encoder = new TextEncoder();
+    const bytes = encoder.encode(text);
+    const len = Math.min(bytes.length, HRM_MAX_LEN);
+
+    const inputIds = new BigInt64Array(HRM_MAX_LEN);
+    const attentionMask = new BigInt64Array(HRM_MAX_LEN);
+
+    for (let i = 0; i < len; i++) {
+      inputIds[i] = BigInt(bytes[i]);
+      attentionMask[i] = 1n;
+    }
+    // Remaining positions are already 0 (padded)
+
+    return { inputIds, attentionMask };
+  }
+
+  async function analyzeDistilbert(text) {
+    const [result] = await distilbertClassifier(text, { topk: 1 });
+    const isInjection = result.label === 'injection' || result.label === 'LABEL_1';
+    return { isInjection, score: result.score };
+  }
+
+  async function analyzeHrm(text) {
+    const { inputIds, attentionMask } = tokenizeBytes(text);
+
+    const inputTensor = new ort.Tensor('int64', inputIds, [1, HRM_MAX_LEN]);
+    const maskTensor = new ort.Tensor('int64', attentionMask, [1, HRM_MAX_LEN]);
+
+    const results = await hrmSession.run({
+      input_ids: inputTensor,
+      attention_mask: maskTensor,
+    });
+
+    const logits = Array.from(results.logits.data);
+    const probs = softmax(logits);
+
+    const isInjection = logits[1] > logits[0];
+    const score = isInjection ? probs[1] : probs[0];
+
+    return { isInjection, score };
+  }
+
+  async function switchModel(model) {
+    if (model === activeModel && !isLoading) return;
+    if (isLoading) return;
+
+    isLoading = true;
+    analyzeBtn.disabled = true;
+    activeModel = model;
+
+    toggleDistilbert.classList.toggle('active', model === 'distilbert');
+    toggleHrm.classList.toggle('active', model === 'hrm');
+
+    updateBadge('Loading...');
+
+    try {
+      if (model === 'distilbert') {
+        if (!distilbertClassifier) {
+          await loadDistilbert();
+        }
+        updateBadge('DistilBERT v3 ready');
+      } else {
+        if (!hrmSession) {
+          await loadHrm();
+        }
+        updateBadge('HRM-Text ready');
+      }
+      hideLoading();
+      analyzeBtn.disabled = false;
+    } catch (err) {
+      updateBadge('Error');
     }
+
+    isLoading = false;
   }
 
   async function analyze() {
     const text = inputEl.value.trim();
-    if (!text || !classifier) return;
+    if (!text) return;
+    if (activeModel === 'distilbert' && !distilbertClassifier) return;
+    if (activeModel === 'hrm' && !hrmSession) return;
 
     analyzeBtn.disabled = true;
     analyzeBtn.textContent = '...';
 
     const start = performance.now();
-    const [result] = await classifier(text, { topk: 1 });
-    const elapsed = Math.round(performance.now() - start);
+    let result;
 
-    const isInjection = result.label === 'injection' || result.label === 'LABEL_1';
-    const score = result.score;
-    const pct = (score * 100).toFixed(1);
+    if (activeModel === 'distilbert') {
+      result = await analyzeDistilbert(text);
+    } else {
+      result = await analyzeHrm(text);
+    }
 
-    resultEl.className = `result visible ${isInjection ? 'danger' : 'safe'}`;
-    resultLabel.textContent = isInjection ? 'Injection Detected' : 'Safe';
+    const elapsed = Math.round(performance.now() - start);
+    const pct = (result.score * 100).toFixed(1);
+
+    resultEl.className = `result visible ${result.isInjection ? 'danger' : 'safe'}`;
+    resultLabel.textContent = result.isInjection ? 'Injection Detected' : 'Safe';
     resultConfidence.textContent = `${pct}%`;
     confidenceFill.style.width = `${pct}%`;
-    latencyEl.textContent = `${elapsed}ms inference`;
+    const modelLabel = activeModel === 'distilbert' ? 'DistilBERT v3' : 'HRM-Text';
+    latencyEl.textContent = `${elapsed}ms inference · ${modelLabel}`;
 
     analyzeBtn.disabled = false;
     analyzeBtn.textContent = 'Analyze';
@@ -433,7 +614,11 @@
     });
   });
 
-  loadModel();
+  toggleDistilbert.addEventListener('click', () => switchModel('distilbert'));
+  toggleHrm.addEventListener('click', () => switchModel('hrm'));
+
+  // Load default model (DistilBERT v3)
+  switchModel('distilbert');
 </script>
 
 </body>