"""SecEmbed interactive retrieval console (ZeroGPU). Pipeline: SOC alert / analyst query → SecEmbed bi-encoder (top-50) → SecReranker cross-encoder (top-5) → ATT&CK / Sigma / CVE / CWE / playbook results Falls back to BGE-small if fine-tuned SecEmbed weights are not yet available. """ from __future__ import annotations import spaces # noqa: F401 — must import before torch import json import os import traceback import gradio as gr import torch from datasets import load_dataset from huggingface_hub import model_info from sentence_transformers import CrossEncoder, SentenceTransformer, util INDEX_REPO = "alirezaaminzadeh/secembed-retrieval-index" SMALL_REPO = "alirezaaminzadeh/SecEmbed-small" BASE_REPO = "alirezaaminzadeh/SecEmbed-base" RERANK_REPO = "alirezaaminzadeh/SecReranker" FALLBACK_EMBED = "BAAI/bge-small-en-v1.5" FALLBACK_RERANK = "cross-encoder/ms-marco-MiniLM-L-6-v2" TOKEN = os.environ.get("HF_TOKEN") def hub_has_weights(repo_id: str) -> bool: try: info = model_info(repo_id, token=TOKEN) files = {s.rfilename for s in (info.siblings or [])} return any( f.endswith(ext) for f in files for ext in (".safetensors", "pytorch_model.bin", "model.safetensors") ) or "config_sentence_transformers.json" in files or "modules.json" in files except Exception: return False def resolve_embed(prefer: str) -> str: order = { "SecEmbed-base": [BASE_REPO, SMALL_REPO, FALLBACK_EMBED], "SecEmbed-small": [SMALL_REPO, FALLBACK_EMBED], "BGE-small baseline": [FALLBACK_EMBED], }[prefer] for rid in order: if rid == FALLBACK_EMBED or hub_has_weights(rid): return rid return FALLBACK_EMBED def resolve_reranker(use_rerank: bool) -> str | None: if not use_rerank: return None if hub_has_weights(RERANK_REPO): return RERANK_REPO return FALLBACK_RERANK print("[info] loading retrieval index", flush=True) try: INDEX = load_dataset(INDEX_REPO, split="train", token=TOKEN) except Exception as e: print(f"[warn] index load failed ({e}); using bootstrap index", flush=True) from datasets import Dataset INDEX = Dataset.from_list([ { "doc_id": "T1059.001", "family": "attack", "title": "T1059.001 PowerShell", "text": "T1059.001 PowerShell. Adversaries may abuse PowerShell commands and scripts for execution.", }, { "doc_id": "T1003.001", "family": "attack", "title": "T1003.001 LSASS Memory", "text": "T1003.001 LSASS Memory. Adversaries may attempt to access credential material stored in process memory of LSASS.", }, { "doc_id": "T1021.001", "family": "attack", "title": "T1021.001 Remote Desktop Protocol", "text": "T1021.001 Remote Desktop Protocol. Adversaries may use Valid Accounts to log into a computer using RDP.", }, { "doc_id": "CWE-89", "family": "cwe", "title": "CWE-89 SQL Injection", "text": "CWE-89 Improper Neutralization of Special Elements used in an SQL Command ('SQL Injection').", }, { "doc_id": "PB-T1059.001", "family": "playbook", "title": "Playbook T1059.001", "text": "Playbook for T1059.001 PowerShell. Validate encoded command telemetry, scope hosts, contain, hunt, eradicate.", }, ]) DOCS = [dict(r) for r in INDEX] DOC_TEXTS = [d["text"] for d in DOCS] print(f"[info] index size={len(DOCS)}", flush=True) # Module-scope models (ZeroGPU registers CUDA tensors at startup) ACTIVE_EMBED_ID = resolve_embed("SecEmbed-small") print(f"[info] embedding model: {ACTIVE_EMBED_ID}", flush=True) EMBED = SentenceTransformer(ACTIVE_EMBED_ID, device="cuda") EMBED.max_seq_length = 256 RERANK_ID = resolve_reranker(True) print(f"[info] reranker model: {RERANK_ID}", flush=True) RERANK = CrossEncoder(RERANK_ID, device="cuda") if RERANK_ID else None # Precompute corpus embeddings for the default embedder print("[info] encoding corpus", flush=True) DOC_EMB = EMBED.encode(DOC_TEXTS, convert_to_tensor=True, show_progress_bar=False) print("[info] ready", flush=True) EXAMPLES = [ "detect powershell encoded command", "remote desktop credential attack", "LSASS memory dump credential dumping", "Suspicious scheduled task creation on Windows endpoint", "CVE related to remote code execution in web servers", "playbook for alert: Encoded PowerShell abuse matching T1059.001", "Sigma rule for suspicious WMI process call", "CWE for SQL injection", ] @spaces.GPU(duration=20) def retrieve(query: str, family: str, top_k_embed: int, top_k_final: int, use_rerank: bool, model_choice: str): global EMBED, DOC_EMB, ACTIVE_EMBED_ID, RERANK, RERANK_ID query = (query or "").strip() if not query: return [], "Enter a SOC alert or analyst query." try: # Hot-swap embedder if user selected a different model that is available wanted = resolve_embed(model_choice) if wanted != ACTIVE_EMBED_ID: ACTIVE_EMBED_ID = wanted EMBED = SentenceTransformer(wanted, device="cuda") EMBED.max_seq_length = 256 DOC_EMB = EMBED.encode(DOC_TEXTS, convert_to_tensor=True, show_progress_bar=False) # Optional family filter if family == "all": idxs = list(range(len(DOCS))) emb = DOC_EMB else: idxs = [i for i, d in enumerate(DOCS) if d["family"] == family] if not idxs: return [], f"No documents for family={family}" emb = DOC_EMB[idxs] q = EMBED.encode(query, convert_to_tensor=True, show_progress_bar=False) scores = util.cos_sim(q, emb)[0] k = min(int(top_k_embed), scores.shape[0]) top = torch.topk(scores, k=k) cand_local = top.indices.tolist() cand_scores = top.values.tolist() cand_global = [idxs[i] for i in cand_local] pairs = [(query, DOCS[i]["text"]) for i in cand_global] if use_rerank: rid = resolve_reranker(True) if rid and rid != RERANK_ID: RERANK_ID = rid RERANK = CrossEncoder(rid, device="cuda") if RERANK is not None: rr = RERANK.predict(pairs) order = sorted(range(len(rr)), key=lambda i: float(rr[i]), reverse=True) cand_global = [cand_global[i] for i in order] cand_scores = [float(rr[i]) for i in order] stage = f"SecEmbed → top-{k} → SecReranker → top-{top_k_final}" else: stage = f"SecEmbed → top-{top_k_final} (reranker unavailable)" else: stage = f"SecEmbed → top-{top_k_final}" final_n = min(int(top_k_final), len(cand_global)) rows = [] for rank, (gi, sc) in enumerate(zip(cand_global[:final_n], cand_scores[:final_n]), start=1): d = DOCS[gi] rows.append([ rank, d["family"], d["doc_id"], d["title"], round(float(sc), 4), d["text"][:420] + ("…" if len(d["text"]) > 420 else ""), ]) meta = ( f"embedder=`{ACTIVE_EMBED_ID}` · reranker=`{RERANK_ID if use_rerank else 'off'}` · " f"corpus={len(DOCS)} · {stage}" ) return rows, meta except Exception: return [], traceback.format_exc() with gr.Blocks(title="SecEmbed") as demo: gr.Markdown( """ # SecEmbed Cybersecurity embedding retrieval for ATT&CK, Sigma, CVE/CWE, and SOC playbooks. **Pipeline:** query → **SecEmbed** (bi-encoder, top-50) → **SecReranker** (cross-encoder, top-5) """ ) with gr.Row(): query = gr.Textbox( label="SOC alert / analyst query", lines=3, placeholder="detect powershell encoded command", ) with gr.Row(): family = gr.Dropdown( choices=["all", "attack", "sigma", "cve", "cwe", "playbook"], value="all", label="Corpus filter", ) model_choice = gr.Dropdown( choices=["SecEmbed-small", "SecEmbed-base", "BGE-small baseline"], value="SecEmbed-small", label="Embedding model", ) use_rerank = gr.Checkbox(value=True, label="Enable SecReranker") top_k_embed = gr.Slider(10, 50, value=50, step=5, label="Bi-encoder top-K") top_k_final = gr.Slider(3, 10, value=5, step=1, label="Final top-K") btn = gr.Button("Retrieve", variant="primary") meta = gr.Markdown() table = gr.Dataframe( headers=["rank", "family", "doc_id", "title", "score", "snippet"], datatype=["number", "str", "str", "str", "number", "str"], label="Results", wrap=True, ) gr.Examples(examples=EXAMPLES, inputs=query) btn.click( retrieve, inputs=[query, family, top_k_embed, top_k_final, use_rerank, model_choice], outputs=[table, meta], ) demo.queue(max_size=8).launch()