Upload precompute.py with huggingface_hub

precompute.py

@@ -0,0 +1,196 @@
+"""
+GuardLLM - Precompute Embeddings & t-SNE
+Downloads the neuralchemy/Prompt-injection-dataset (core config),
+extracts CLS embeddings from Llama Prompt Guard 2 (86M),
+computes t-SNE 2D projection, and saves everything to a cache file.
+
+Run this script ONCE before launching the app (or let the app run it on first start).
+"""
+
+import os
+import json
+import logging
+import numpy as np
+import torch
+from pathlib import Path
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger("precompute")
+
+CACHE_DIR = Path(__file__).parent / "cache"
+CACHE_FILE = CACHE_DIR / "embeddings_tsne.npz"
+META_FILE = CACHE_DIR / "metadata.json"
+MODEL_ID = "meta-llama/Llama-Prompt-Guard-2-86M"
+DATASET_ID = "neuralchemy/Prompt-injection-dataset"
+DATASET_CONFIG = "core"
+BATCH_SIZE = 32
+MAX_LENGTH = 512
+TSNE_PERPLEXITY = 30
+TSNE_SEED = 42
+
+
+def is_cached() -> bool:
+    """Check if precomputed data exists."""
+    return CACHE_FILE.exists() and META_FILE.exists()
+
+
+def load_cached():
+    """Load precomputed embeddings, t-SNE coords, and metadata."""
+    logger.info("Loading cached data from %s", CACHE_DIR)
+    data = np.load(CACHE_FILE)
+    with open(META_FILE, "r", encoding="utf-8") as f:
+        metadata = json.load(f)
+    return {
+        "embeddings": data["embeddings"],
+        "tsne_2d": data["tsne_2d"],
+        "metadata": metadata,
+    }
+
+
+def download_dataset():
+    """Download the neuralchemy dataset (core config)."""
+    from datasets import load_dataset
+
+    logger.info("Downloading dataset %s (config=%s)...", DATASET_ID, DATASET_CONFIG)
+    ds = load_dataset(DATASET_ID, DATASET_CONFIG)
+
+    # Combine all splits for the visualization
+    all_samples = []
+    for split_name in ["train", "validation", "test"]:
+        if split_name in ds:
+            split = ds[split_name]
+            logger.info("  Split '%s': %d samples", split_name, len(split))
+            for row in split:
+                all_samples.append({
+                    "text": row["text"],
+                    "label": int(row["label"]),
+                    "category": row.get("category", "unknown"),
+                    "severity": row.get("severity", ""),
+                    "source": row.get("source", ""),
+                    "split": split_name,
+                })
+
+    logger.info("Total samples: %d", len(all_samples))
+    return all_samples
+
+
+def compute_embeddings(samples: list) -> np.ndarray:
+    """Extract CLS token embeddings from Llama Prompt Guard 2."""
+    from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+    logger.info("Loading model %s...", MODEL_ID)
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        MODEL_ID, output_hidden_states=True
+    )
+    model.eval()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    logger.info("Using device: %s", device)
+
+    texts = [s["text"] for s in samples]
+    all_embeddings = []
+
+    num_batches = (len(texts) + BATCH_SIZE - 1) // BATCH_SIZE
+    for i in range(0, len(texts), BATCH_SIZE):
+        batch_idx = i // BATCH_SIZE + 1
+        batch_texts = texts[i : i + BATCH_SIZE]
+
+        if batch_idx % 10 == 1 or batch_idx == num_batches:
+            logger.info("  Batch %d/%d (%d samples)...", batch_idx, num_batches, len(batch_texts))
+
+        inputs = tokenizer(
+            batch_texts,
+            return_tensors="pt",
+            truncation=True,
+            max_length=MAX_LENGTH,
+            padding=True,
+        ).to(device)
+
+        with torch.no_grad():
+            outputs = model(**inputs)
+            # CLS token embedding from last hidden layer
+            hidden_states = outputs.hidden_states[-1]  # [batch, seq_len, 768]
+            cls_embeddings = hidden_states[:, 0, :].cpu().numpy()  # [batch, 768]
+            all_embeddings.append(cls_embeddings)
+
+    embeddings = np.concatenate(all_embeddings, axis=0)
+    logger.info("Embeddings shape: %s", embeddings.shape)
+    return embeddings
+
+
+def compute_tsne(embeddings: np.ndarray) -> np.ndarray:
+    """Run t-SNE dimensionality reduction to 2D."""
+    from sklearn.manifold import TSNE
+
+    n_samples = embeddings.shape[0]
+    perplexity = min(TSNE_PERPLEXITY, n_samples - 1)
+
+    logger.info(
+        "Running t-SNE (n=%d, perplexity=%d, random_state=%d)...",
+        n_samples, perplexity, TSNE_SEED,
+    )
+    tsne = TSNE(
+        n_components=2,
+        perplexity=perplexity,
+        random_state=TSNE_SEED,
+        n_iter=1000,
+        learning_rate="auto",
+        init="pca",
+    )
+    coords_2d = tsne.fit_transform(embeddings)
+    logger.info("t-SNE done. Output shape: %s", coords_2d.shape)
+    return coords_2d
+
+
+def precompute_all():
+    """Full pipeline: download → embed → t-SNE → save."""
+    if is_cached():
+        logger.info("Cache already exists. Loading...")
+        return load_cached()
+
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+    # Step 1: Download dataset
+    samples = download_dataset()
+
+    # Step 2: Compute embeddings
+    embeddings = compute_embeddings(samples)
+
+    # Step 3: Compute t-SNE
+    tsne_2d = compute_tsne(embeddings)
+
+    # Step 4: Save
+    logger.info("Saving to cache...")
+    np.savez_compressed(
+        CACHE_FILE,
+        embeddings=embeddings,
+        tsne_2d=tsne_2d,
+    )
+
+    # Save metadata (text, labels, categories) as JSON
+    metadata = []
+    for s in samples:
+        metadata.append({
+            "text": s["text"],
+            "label": s["label"],
+            "category": s["category"],
+            "severity": s["severity"],
+            "source": s["source"],
+            "split": s["split"],
+        })
+
+    with open(META_FILE, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, ensure_ascii=False)
+
+    logger.info("Cache saved to %s", CACHE_DIR)
+    return {
+        "embeddings": embeddings,
+        "tsne_2d": tsne_2d,
+        "metadata": metadata,
+    }
+
+
+if __name__ == "__main__":
+    precompute_all()