Upload app.py with huggingface_hub

app.py

@@ -14,9 +14,7 @@ import torch
 import numpy as np
 import plotly.graph_objects as go
 import plotly.io as pio
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
-from precompute import precompute_all, is_cached
+from pathlib import Path
 
 # ---------------------------------------------------------------------------
 # Logging
@@ -64,29 +62,49 @@ CATEGORY_LABELS_FR = {
 }
 
 # ---------------------------------------------------------------------------
-# Load model for real-time analysis
+# Lazy-loaded risk classifier (Llama Prompt Guard 2)
+# Loaded on first use so Space startup is fast.
 # ---------------------------------------------------------------------------
 MODEL_ID = "meta-llama/Llama-Prompt-Guard-2-86M"
+LABELS = ["Benign", "Malicious"]
 
-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("Model loaded on %s", DEVICE)
+_classifier = {"tokenizer": None, "model": None, "device": None}
+
+
+def get_classifier():
+    if _classifier["model"] is None:
+        logger.info("Lazy-loading Llama Prompt Guard 2...")
+        from transformers import AutoTokenizer, AutoModelForSequenceClassification
+        tok = AutoTokenizer.from_pretrained(MODEL_ID)
+        mdl = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
+        mdl.eval()
+        dev = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        mdl.to(dev)
+        _classifier["tokenizer"] = tok
+        _classifier["model"] = mdl
+        _classifier["device"] = dev
+        logger.info("Classifier loaded on %s", dev)
+    return _classifier["tokenizer"], _classifier["model"], _classifier["device"]
 
-LABELS = ["Benign", "Malicious"]
 
 # ---------------------------------------------------------------------------
-# Load precomputed t-SNE data
+# Load precomputed t-SNE data (fail fast if cache missing)
 # ---------------------------------------------------------------------------
-logger.info("Loading precomputed embeddings & t-SNE...")
-cached_data = precompute_all()
-TSNE_COORDS = cached_data["tsne_2d"]
-METADATA = cached_data["metadata"]
+CACHE_DIR = Path(__file__).parent / "cache"
+CACHE_FILE = CACHE_DIR / "embeddings_tsne.npz"
+META_FILE = CACHE_DIR / "metadata.json"
+
+logger.info("Loading precomputed t-SNE cache from %s", CACHE_DIR)
+if not CACHE_FILE.exists() or not META_FILE.exists():
+    raise RuntimeError(
+        f"Cache files not found in {CACHE_DIR}. "
+        f"Run `python precompute.py` first (or ensure cache/ is committed to the Space)."
+    )
+
+_npz = np.load(CACHE_FILE)
+TSNE_COORDS = _npz["tsne_2d"]
+with open(META_FILE, "r", encoding="utf-8") as f:
+    METADATA = json.load(f)
 logger.info("Loaded %d points for visualization", len(METADATA))
 
 ALL_TEXTS = [m["text"] for m in METADATA]
@@ -108,10 +126,11 @@ for i, m in enumerate(METADATA):
 # Analysis function
 # ---------------------------------------------------------------------------
 def analyze_prompt(text: str):
-    """Run Llama Prompt Guard 2 on a single prompt."""
+    """Run Llama Prompt Guard 2 on a single prompt (lazy-loads on first call)."""
     if not text or not text.strip():
         return {}, 0.0
 
+    tokenizer, model, DEVICE = get_classifier()
     inputs = tokenizer(
         text, return_tensors="pt", truncation=True, max_length=512, padding=True
     ).to(DEVICE)
@@ -618,31 +637,4 @@ with gr.Blocks(
         inputs=[manual_input],
         outputs=[result_html, risk_md],
     )
-    manual_input.submit(
-        fn=on_manual_analyze,
-        inputs=[manual_input],
-        outputs=[result_html, risk_md],
-    )
-
-    # Inject Plotly click handler JS
-    demo.load(fn=None, inputs=None, outputs=None, js=PLOTLY_CLICK_JS)
-
-    # Footer
-    gr.Markdown(
-        """
-        ---
-        <div style="text-align:center; color:#64748b; font-size:0.8em;">
-            <strong>GuardLLM</strong> — Visualiseur de sécurité des prompts<br>
-            Modèle : <a href="https://huggingface.co/meta-llama/Llama-Prompt-Guard-2-86M">
-            Llama Prompt Guard 2 (86M)</a> par Meta &bull;
-            Dataset : <a href="https://huggingface.co/datasets/neuralchemy/Prompt-injection-dataset">
-            neuralchemy/Prompt-injection-dataset</a>
-        </div>
-        """
-    )
-
-
-logger.info("Gradio app built. Ready to launch.")
-
-if __name__ == "__main__":
-    demo.launch()
+    manua