"""SecEmbed dataset builder. Constructs contrastive training triplets and a multi-task cybersecurity retrieval benchmark from public security corpora: - MITRE ATT&CK Enterprise techniques (STIX 2.1) - SigmaHQ detection rules (via SigmaForge Hub mirror + SigmaHQ raw) - CVE descriptions (CVE-RiskRank Hub dataset) - CWE weakness descriptions (MITRE CWE XML/CSV mirror) - Curated SOC-alert and threat-intel query templates All downloads and processing run on Hugging Face infrastructure. Outputs are pushed to: - alirezaaminzadeh/secembed-pairs - alirezaaminzadeh/cybersec-retrieval-benchmark """ from __future__ import annotations import hashlib import json import os import random import re import zipfile from io import BytesIO from typing import Any import requests from datasets import Dataset, DatasetDict, Features, Sequence, Value, load_dataset TOKEN = os.environ.get("HF_TOKEN") PAIRS_REPO = os.environ.get("PAIRS_REPO", "alirezaaminzadeh/secembed-pairs") BENCH_REPO = os.environ.get("BENCH_REPO", "alirezaaminzadeh/cybersec-retrieval-benchmark") SEED = 20260811 rng = random.Random(SEED) ATTACK_URL = ( "https://raw.githubusercontent.com/mitre/cti/master/enterprise-attack/" "enterprise-attack.json" ) CWE_URL = "https://cwe.mitre.org/data/csv/1000.csv.zip" SIGMA_MIRROR = "alirezaaminzadeh/sigmaforge-detection-rules" CVE_MIRROR = "alirezaaminzadeh/cve-riskrank" UA = {"User-Agent": "secembed-builder/1.0"} # --------------------------------------------------------------------------- # Query templates (analyst / SOC / TI language) # --------------------------------------------------------------------------- ATTACK_QUERY_TEMPLATES = [ "detect {name}", "how to detect {name}", "ATT&CK technique for {name}", "indicators of {name}", "hunt for {name} activity", "SOC alert related to {name}", "adversary uses {name}", "mitigate {name}", "{tactic} using {name}", "detection engineering for {name}", "what is {tid}", "explain {tid} {name}", "telemetry for {tid}", ] SIGMA_QUERY_TEMPLATES = [ "Sigma rule for {title}", "detection rule: {title}", "write Sigma for {title}", "alert when {title}", "detect {title} in logs", "rule covering {title}", "Sigma YAML for {logsource}", ] CVE_QUERY_TEMPLATES = [ "vulnerability: {summary}", "CVE similar to {cve_id}", "exploit {summary}", "patch priority for {cve_id}", "describe {cve_id}", "risk of {summary}", ] CWE_QUERY_TEMPLATES = [ "weakness {name}", "CWE for {name}", "explain {cwe_id}", "common weakness: {name}", "secure coding against {name}", ] THREAT_REPORT_TEMPLATES = [ "threat report mentioning {name}", "campaign using {name}", "actor TTPs: {name}", "intelligence brief on {tid}", ] PLAYBOOK_QUERY_TEMPLATES = [ "playbook for alert: {alert}", "incident response steps for {alert}", "SOC runbook when seeing {alert}", "containment actions for {alert}", ] SOC_ALERT_TEMPLATES = [ "Suspicious {name} execution detected on endpoint", "EDR alert: possible {name} ({tid})", "SIEM correlation: {tactic} activity consistent with {name}", "Multiple failed authentications followed by {name}", "Outbound C2 pattern associated with {name}", "Encoded PowerShell / command-line abuse matching {name}", "Lateral movement indicators for {name}", "Persistence mechanism resembling {name}", ] def http_get(url: str, timeout: int = 180) -> bytes: last: Exception | None = None for attempt in range(4): try: r = requests.get(url, timeout=timeout, headers=UA) r.raise_for_status() return r.content except Exception as e: # noqa: BLE001 last = e print(f"[warn] GET {url} attempt {attempt + 1}: {e}", flush=True) raise RuntimeError(f"failed to download {url}: {last}") def clean(text: str, limit: int = 1200) -> str: text = re.sub(r"\s+", " ", (text or "")).strip() return text[:limit] def uid(*parts: str) -> str: h = hashlib.sha1("|".join(parts).encode("utf-8")).hexdigest()[:16] return h # --------------------------------------------------------------------------- # Loaders # --------------------------------------------------------------------------- def load_attack() -> list[dict[str, Any]]: print("[info] loading MITRE ATT&CK Enterprise", flush=True) data = json.loads(http_get(ATTACK_URL)) techniques: list[dict[str, Any]] = [] for obj in data.get("objects", []): if obj.get("type") != "attack-pattern" or obj.get("revoked"): continue ext = obj.get("external_references") or [] tid = next((e.get("external_id") for e in ext if e.get("source_name") == "mitre-attack"), None) if not tid or not tid.startswith("T"): continue tactics = [] for phase in obj.get("kill_chain_phases") or []: if phase.get("kill_chain_name") == "mitre-attack": tactics.append(phase.get("phase_name", "").replace("-", " ")) name = obj.get("name", "") desc = clean(obj.get("description", ""), 2000) detection = clean(obj.get("x_mitre_detection", ""), 1500) passage = ( f"{tid} {name}. Tactics: {', '.join(tactics) or 'n/a'}. " f"Description: {desc}" ) if detection: passage += f" Detection: {detection}" techniques.append({ "id": tid, "name": name, "tactics": tactics, "description": desc, "detection": detection, "passage": clean(passage, 2500), "parent": tid.split(".")[0] if "." in tid else tid, }) print(f"[info] ATT&CK techniques: {len(techniques)}", flush=True) return techniques def load_sigma() -> list[dict[str, Any]]: print("[info] loading Sigma rules from Hub mirror", flush=True) rules: list[dict[str, Any]] = [] try: ds = load_dataset(SIGMA_MIRROR, "description_to_sigma", split="train", token=TOKEN) for row in ds: title = clean(row.get("description", "")[:120] or "Sigma rule", 120) body = clean(row.get("sigma_rule", ""), 2500) if len(body) < 40: continue ls = " / ".join( filter( None, [ row.get("logsource_product") or "", row.get("logsource_category") or "", row.get("logsource_service") or "", ], ) ) attack = row.get("attack_techniques") or [] if isinstance(attack, str): attack = [attack] passage = f"Sigma rule: {title}. Logsource: {ls or 'n/a'}. Rule:\n{body}" rules.append({ "id": uid("sigma", title, body[:80]), "title": title, "logsource": ls, "attack": list(attack), "passage": clean(passage, 2800), }) except Exception as e: # noqa: BLE001 print(f"[warn] Sigma Hub load failed: {e}", flush=True) if len(rules) < 200: print("[info] supplementing with synthetic Sigma-style passages", flush=True) extras = [ ("PowerShell Encoded Command", "windows", "process_creation", ["T1059.001"]), ("Suspicious WMI Process Call", "windows", "wmi_event", ["T1047"]), ("Remote Desktop Logon", "windows", "security", ["T1021.001"]), ("Scheduled Task Creation", "windows", "process_creation", ["T1053.005"]), ("LSASS Memory Access", "windows", "process_access", ["T1003.001"]), ("CertUtil Download", "windows", "process_creation", ["T1105"]), ("Suspicious Service Installation", "windows", "system", ["T1543.003"]), ("DNS Tunneling Pattern", "linux", "network_connection", ["T1071.004"]), ("SSH Brute Force", "linux", "auth", ["T1110"]), ("Kubernetes Privileged Pod", "kubernetes", "audit", ["T1611"]), ] for title, product, category, attack in extras: yaml_body = ( f"title: {title}\nstatus: experimental\nlogsource:\n" f" product: {product}\n category: {category}\ndetection:\n" f" selection:\n CommandLine|contains: '{title.split()[0]}'\n" f" condition: selection\nlevel: high\n" ) rules.append({ "id": uid("sigma-syn", title), "title": title, "logsource": f"{product}/{category}", "attack": attack, "passage": clean( f"Sigma rule: {title}. Logsource: {product}/{category}. Rule:\n{yaml_body}", 2800, ), }) print(f"[info] Sigma rules: {len(rules)}", flush=True) return rules def load_cves(limit: int = 8000) -> list[dict[str, Any]]: print("[info] loading CVE descriptions from Hub", flush=True) cves: list[dict[str, Any]] = [] try: ds = load_dataset(CVE_MIRROR, split="train", token=TOKEN) for i, row in enumerate(ds): if i >= limit: break cve_id = row.get("cve_id") or "" desc = clean(row.get("description", ""), 1800) if not cve_id or len(desc) < 40: continue cwe = row.get("primary_cwe") or "UNKNOWN" score = row.get("cvss_base_score") passage = f"{cve_id}. CWE: {cwe}. CVSS: {score}. {desc}" cves.append({ "id": cve_id, "cwe": str(cwe), "summary": desc[:180], "passage": clean(passage, 2200), }) except Exception as e: # noqa: BLE001 print(f"[warn] CVE Hub load failed: {e}", flush=True) print(f"[info] CVEs: {len(cves)}", flush=True) return cves def load_cwes() -> list[dict[str, Any]]: print("[info] loading CWE catalog", flush=True) cwes: list[dict[str, Any]] = [] try: raw = http_get(CWE_URL) with zipfile.ZipFile(BytesIO(raw)) as zf: name = next(n for n in zf.namelist() if n.endswith(".csv")) text = zf.read(name).decode("utf-8", errors="ignore") # CSV has quoted fields; use a light parser import csv reader = csv.DictReader(text.splitlines()) for row in reader: cid = (row.get("CWE-ID") or row.get("ID") or "").strip() name = (row.get("Name") or "").strip() desc = clean(row.get("Description") or row.get("Extended Description") or "", 1800) if not cid or not name: continue cwe_id = f"CWE-{cid}" if not cid.startswith("CWE-") else cid passage = f"{cwe_id} {name}. {desc}" cwes.append({ "id": cwe_id, "name": name, "passage": clean(passage, 2200), }) except Exception as e: # noqa: BLE001 print(f"[warn] CWE download failed ({e}); using seed list", flush=True) seed = [ ("CWE-79", "Cross-site Scripting", "Improper neutralization of input during web page generation."), ("CWE-89", "SQL Injection", "Improper neutralization of special elements used in an SQL command."), ("CWE-22", "Path Traversal", "Improper limitation of a pathname to a restricted directory."), ("CWE-287", "Improper Authentication", "When an actor claims to have a given identity, the software does not prove or insufficiently proves that claim."), ("CWE-502", "Deserialization of Untrusted Data", "The application deserializes untrusted data without sufficiently verifying the resulting data will be valid."), ("CWE-798", "Hard-coded Credentials", "The software contains hard-coded credentials, such as a password or cryptographic key."), ("CWE-119", "Buffer Overflow", "The software performs operations on a memory buffer, but it can read from or write to a memory location outside of the intended boundary."), ("CWE-352", "CSRF", "The web application does not verify that a request was intentionally provided by the user who submitted it."), ] for cid, name, desc in seed: cwes.append({"id": cid, "name": name, "passage": f"{cid} {name}. {desc}"}) print(f"[info] CWEs: {len(cwes)}", flush=True) return cwes # --------------------------------------------------------------------------- # Pair / hard-negative construction # --------------------------------------------------------------------------- def fmt(template: str, **kwargs: str) -> str: try: return template.format(**kwargs) except KeyError: return template def hard_neg_attack(tech: dict, pool: list[dict], k: int = 1) -> list[dict]: parent = tech["parent"] same_parent = [t for t in pool if t["parent"] == parent and t["id"] != tech["id"]] same_tactic = [ t for t in pool if t["id"] != tech["id"] and set(t["tactics"]) & set(tech["tactics"]) ] candidates = same_parent or same_tactic or [t for t in pool if t["id"] != tech["id"]] return rng.sample(candidates, min(k, len(candidates))) def hard_neg_other(item: dict, pool: list[dict], id_key: str = "id", k: int = 1) -> list[dict]: others = [x for x in pool if x[id_key] != item[id_key]] if not others: return [] return rng.sample(others, min(k, len(others))) def build_triplets( techniques: list[dict], rules: list[dict], cves: list[dict], cwes: list[dict], ) -> list[dict]: rows: list[dict] = [] # ATT&CK for tech in techniques: queries = [ fmt(t, name=tech["name"], tid=tech["id"], tactic=(tech["tactics"][0] if tech["tactics"] else "adversary")) for t in ATTACK_QUERY_TEMPLATES ] queries += [ fmt(t, name=tech["name"], tid=tech["id"]) for t in THREAT_REPORT_TEMPLATES ] queries += [ fmt(t, name=tech["name"], tid=tech["id"], tactic=(tech["tactics"][0] if tech["tactics"] else "tactic")) for t in SOC_ALERT_TEMPLATES ] for q in queries: for neg in hard_neg_attack(tech, techniques, k=1): rows.append({ "pair_id": uid("attack", tech["id"], q), "task": "attack_retrieval", "query": q, "positive": tech["passage"], "hard_negative": neg["passage"], "positive_id": tech["id"], "negative_id": neg["id"], "source": "mitre-attack", }) # Sigma for rule in rules: queries = [ fmt(t, title=rule["title"], logsource=rule["logsource"] or "windows") for t in SIGMA_QUERY_TEMPLATES ] for q in queries: for neg in hard_neg_other(rule, rules, k=1): rows.append({ "pair_id": uid("sigma", rule["id"], q), "task": "sigma_retrieval", "query": q, "positive": rule["passage"], "hard_negative": neg["passage"], "positive_id": rule["id"], "negative_id": neg["id"], "source": "sigmahq", }) # CVE for cve in cves: queries = [fmt(t, summary=cve["summary"], cve_id=cve["id"]) for t in CVE_QUERY_TEMPLATES] # Prefer hard neg with same CWE when possible same_cwe = [c for c in cves if c["cwe"] == cve["cwe"] and c["id"] != cve["id"]] negs = rng.sample(same_cwe, 1) if same_cwe else hard_neg_other(cve, cves, k=1) for q in queries: for neg in negs: rows.append({ "pair_id": uid("cve", cve["id"], q), "task": "cve_similarity", "query": q, "positive": cve["passage"], "hard_negative": neg["passage"], "positive_id": cve["id"], "negative_id": neg["id"], "source": "nvd", }) # CWE for cwe in cwes: queries = [fmt(t, name=cwe["name"], cwe_id=cwe["id"]) for t in CWE_QUERY_TEMPLATES] for q in queries: for neg in hard_neg_other(cwe, cwes, k=1): rows.append({ "pair_id": uid("cwe", cwe["id"], q), "task": "cwe_retrieval", "query": q, "positive": cwe["passage"], "hard_negative": neg["passage"], "positive_id": cwe["id"], "negative_id": neg["id"], "source": "mitre-cwe", }) # SOC alert → playbook (synthetic playbooks grounded in ATT&CK) for tech in techniques: if rng.random() > 0.55: continue alert = fmt( rng.choice(SOC_ALERT_TEMPLATES), name=tech["name"], tid=tech["id"], tactic=(tech["tactics"][0] if tech["tactics"] else "adversary"), ) playbook = ( f"Playbook for {tech['id']} {tech['name']}. " f"1) Validate alert telemetry against {tech['id']}. " f"2) Scope affected hosts and accounts. " f"3) Contain by isolating endpoint / disabling account if confirmed. " f"4) Hunt for related tactics: {', '.join(tech['tactics']) or 'n/a'}. " f"5) Eradicate persistence and recover. " f"Reference detection notes: {tech['detection'] or tech['description'][:400]}" ) neg_tech = hard_neg_attack(tech, techniques, k=1)[0] neg_playbook = ( f"Playbook for {neg_tech['id']} {neg_tech['name']}. " f"Different response path focused on {', '.join(neg_tech['tactics']) or 'other tactics'}." ) for t in PLAYBOOK_QUERY_TEMPLATES: q = fmt(t, alert=alert) rows.append({ "pair_id": uid("playbook", tech["id"], q), "task": "soc_playbook", "query": q, "positive": clean(playbook, 2200), "hard_negative": clean(neg_playbook, 2200), "positive_id": f"PB-{tech['id']}", "negative_id": f"PB-{neg_tech['id']}", "source": "synthetic-playbook", }) rng.shuffle(rows) print(f"[info] total triplets before dedupe: {len(rows)}", flush=True) seen = set() deduped = [] for r in rows: key = (r["query"], r["positive_id"], r["task"]) if key in seen: continue seen.add(key) deduped.append(r) print(f"[info] triplets after dedupe: {len(deduped)}", flush=True) return deduped def split_rows(rows: list[dict], train_ratio: float = 0.9, val_ratio: float = 0.05) -> DatasetDict: n = len(rows) n_train = int(n * train_ratio) n_val = int(n * val_ratio) train = rows[:n_train] val = rows[n_train:n_train + n_val] test = rows[n_train + n_val:] features = Features({ "pair_id": Value("string"), "task": Value("string"), "query": Value("string"), "positive": Value("string"), "hard_negative": Value("string"), "positive_id": Value("string"), "negative_id": Value("string"), "source": Value("string"), }) return DatasetDict({ "train": Dataset.from_list(train, features=features), "validation": Dataset.from_list(val, features=features), "test": Dataset.from_list(test, features=features), }) # --------------------------------------------------------------------------- # Benchmark # --------------------------------------------------------------------------- def build_benchmark( techniques: list[dict], rules: list[dict], cves: list[dict], cwes: list[dict], ) -> DatasetDict: """Build retrieval tasks: queries + corpus + qrels.""" corpora: dict[str, list[dict]] = { "attack_retrieval": [ {"doc_id": t["id"], "text": t["passage"], "meta": json.dumps({"name": t["name"], "tactics": t["tactics"]})} for t in techniques ], "sigma_retrieval": [ {"doc_id": r["id"], "text": r["passage"], "meta": json.dumps({"title": r["title"], "attack": r["attack"]})} for r in rules ], "cve_similarity": [ {"doc_id": c["id"], "text": c["passage"], "meta": json.dumps({"cwe": c["cwe"]})} for c in cves[:3000] ], "cwe_retrieval": [ {"doc_id": c["id"], "text": c["passage"], "meta": json.dumps({"name": c["name"]})} for c in cwes ], } # Threat report corpus = ATT&CK passages styled as intel blurbs corpora["threat_report_retrieval"] = [ { "doc_id": f"TR-{t['id']}", "text": clean( f"Threat intelligence note: adversaries leveraging {t['id']} ({t['name']}) " f"during {', '.join(t['tactics']) or 'intrusions'}. {t['description'][:900]}", 2000, ), "meta": json.dumps({"technique": t["id"]}), } for t in techniques ] # Playbooks corpora["soc_playbook"] = [ { "doc_id": f"PB-{t['id']}", "text": clean( f"Playbook for {t['id']} {t['name']}. Validate telemetry, scope hosts, " f"contain, hunt related tactics ({', '.join(t['tactics']) or 'n/a'}), eradicate. " f"{t['detection'] or t['description'][:500]}", 2000, ), "meta": json.dumps({"technique": t["id"]}), } for t in techniques ] queries: list[dict] = [] qrels: list[dict] = [] # Sample evaluation queries per task sample_techs = rng.sample(techniques, min(250, len(techniques))) for t in sample_techs: q = fmt( rng.choice(ATTACK_QUERY_TEMPLATES + SOC_ALERT_TEMPLATES), name=t["name"], tid=t["id"], tactic=(t["tactics"][0] if t["tactics"] else "adversary"), ) qid = uid("bq-attack", t["id"], q) queries.append({"query_id": qid, "task": "attack_retrieval", "query": q}) qrels.append({"query_id": qid, "doc_id": t["id"], "relevance": 1, "task": "attack_retrieval"}) q2 = fmt(rng.choice(THREAT_REPORT_TEMPLATES), name=t["name"], tid=t["id"]) qid2 = uid("bq-tr", t["id"], q2) queries.append({"query_id": qid2, "task": "threat_report_retrieval", "query": q2}) qrels.append({"query_id": qid2, "doc_id": f"TR-{t['id']}", "relevance": 1, "task": "threat_report_retrieval"}) alert = fmt( rng.choice(SOC_ALERT_TEMPLATES), name=t["name"], tid=t["id"], tactic=(t["tactics"][0] if t["tactics"] else "adversary"), ) q3 = fmt(rng.choice(PLAYBOOK_QUERY_TEMPLATES), alert=alert) qid3 = uid("bq-pb", t["id"], q3) queries.append({"query_id": qid3, "task": "soc_playbook", "query": q3}) qrels.append({"query_id": qid3, "doc_id": f"PB-{t['id']}", "relevance": 1, "task": "soc_playbook"}) sample_rules = rng.sample(rules, min(200, len(rules))) for r in sample_rules: q = fmt(rng.choice(SIGMA_QUERY_TEMPLATES), title=r["title"], logsource=r["logsource"] or "windows") qid = uid("bq-sigma", r["id"], q) queries.append({"query_id": qid, "task": "sigma_retrieval", "query": q}) qrels.append({"query_id": qid, "doc_id": r["id"], "relevance": 1, "task": "sigma_retrieval"}) sample_cves = rng.sample(cves[:3000], min(250, len(cves[:3000]))) if cves else [] for c in sample_cves: q = fmt(rng.choice(CVE_QUERY_TEMPLATES), summary=c["summary"], cve_id=c["id"]) qid = uid("bq-cve", c["id"], q) queries.append({"query_id": qid, "task": "cve_similarity", "query": q}) qrels.append({"query_id": qid, "doc_id": c["id"], "relevance": 1, "task": "cve_similarity"}) sample_cwes = rng.sample(cwes, min(150, len(cwes))) for c in sample_cwes: q = fmt(rng.choice(CWE_QUERY_TEMPLATES), name=c["name"], cwe_id=c["id"]) qid = uid("bq-cwe", c["id"], q) queries.append({"query_id": qid, "task": "cwe_retrieval", "query": q}) qrels.append({"query_id": qid, "doc_id": c["id"], "relevance": 1, "task": "cwe_retrieval"}) # Flatten corpus with task tag corpus_rows = [] for task, docs in corpora.items(): for d in docs: corpus_rows.append({ "task": task, "doc_id": d["doc_id"], "text": d["text"], "meta": d["meta"], }) print(f"[info] benchmark queries={len(queries)} qrels={len(qrels)} corpus={len(corpus_rows)}", flush=True) return DatasetDict({ "queries": Dataset.from_list(queries, features=Features({ "query_id": Value("string"), "task": Value("string"), "query": Value("string"), })), "corpus": Dataset.from_list(corpus_rows, features=Features({ "task": Value("string"), "doc_id": Value("string"), "text": Value("string"), "meta": Value("string"), })), "qrels": Dataset.from_list(qrels, features=Features({ "query_id": Value("string"), "doc_id": Value("string"), "relevance": Value("int64"), "task": Value("string"), })), }) PAIRS_CARD = """--- license: mit task_categories: - sentence-similarity - feature-extraction language: - en tags: - cybersecurity - retrieval - embedding - mitre-attack - sigma - cve - cwe - soc - contrastive-learning - sentence-transformers pretty_name: SecEmbed Training Pairs size_categories: - 10K None: print(f"[info] pushing {repo_id}", flush=True) if configs: # Heterogeneous splits → one Hub config per split name for name, split in ds.items(): print(f"[info] config={name} rows={len(split)}", flush=True) split.push_to_hub(repo_id, config_name=name, token=TOKEN, private=False) else: ds.push_to_hub(repo_id, token=TOKEN, private=False) from huggingface_hub import HfApi HfApi(token=TOKEN).upload_file( path_or_fileobj=card.encode("utf-8"), path_in_repo="README.md", repo_id=repo_id, repo_type="dataset", commit_message="Add dataset card", ) print(f"[info] published https://huggingface.co/datasets/{repo_id}", flush=True) def main() -> None: print("[info] SecEmbed dataset build starting", flush=True) techniques = load_attack() rules = load_sigma() cves = load_cves() cwes = load_cwes() triplets = build_triplets(techniques, rules, cves, cwes) pairs = split_rows(triplets) for split, d in pairs.items(): print(f"[info] pairs/{split}: {len(d)}", flush=True) push_dataset(pairs, PAIRS_REPO, PAIRS_CARD) bench = build_benchmark(techniques, rules, cves, cwes) for split, d in bench.items(): print(f"[info] bench/{split}: {len(d)}", flush=True) push_dataset(bench, BENCH_REPO, BENCH_CARD, configs=True) # Also publish a compact retrieval index used by the demo Space index_docs = [] for t in techniques: index_docs.append({ "doc_id": t["id"], "family": "attack", "title": f"{t['id']} {t['name']}", "text": t["passage"], }) for r in rules[:2500]: index_docs.append({ "doc_id": r["id"], "family": "sigma", "title": r["title"], "text": r["passage"], }) for c in cves[:2500]: index_docs.append({ "doc_id": c["id"], "family": "cve", "title": c["id"], "text": c["passage"], }) for c in cwes: index_docs.append({ "doc_id": c["id"], "family": "cwe", "title": f"{c['id']} {c['name']}", "text": c["passage"], }) for t in techniques: index_docs.append({ "doc_id": f"PB-{t['id']}", "family": "playbook", "title": f"Playbook {t['id']}", "text": clean( f"Playbook for {t['id']} {t['name']}. Validate, scope, contain, hunt, eradicate. " f"{t['detection'] or t['description'][:500]}", 2000, ), }) index_ds = Dataset.from_list(index_docs) index_repo = "alirezaaminzadeh/secembed-retrieval-index" print(f"[info] pushing retrieval index ({len(index_docs)} docs) → {index_repo}", flush=True) index_ds.push_to_hub(index_repo, token=TOKEN, private=False) print("[info] SecEmbed dataset build complete", flush=True) if __name__ == "__main__": main()