new tests and tests results

Author: Mec-iS

tests/embeddings_model_tsdae.py

@@ -0,0 +1,59 @@
+from pathlib import Path
+from tqdm import tqdm
+
+from tsdae import TSDAE
+import nltk
+
+from test_2_CVE_db import iter_cve_json, extract_text
+
+# --- NLTK setup (reuse your custom dir if you want) ---
+nltk_data_dir = Path(__file__).parent.parent / ".venv" / "nltk_data"
+nltk_data_dir.mkdir(parents=True, exist_ok=True)
+
+nltk.data.path.insert(0, str(nltk_data_dir))
+try:
+    nltk.data.find("tokenizers/punkt_tab/english")
+except LookupError:
+    nltk.download("punkt_tab", download_dir=str(nltk_data_dir))
+
+# --- Load CVE corpus ---
+dataset_root = (
+    Path(__file__).parent.parent.parent.parent / "datasets/cvelistV5-main"
+)
+print(f"Loading dataset at {dataset_root}")
+
+ids, corpus = [], []
+print("Start JSON iteration")
+for _, j in tqdm(iter_cve_json(dataset_root, 2013, 2018)):
+    cve_id, title, text = extract_text(j)
+    ids.append(cve_id)
+    corpus.append(title + "\n" + text)
+
+if not corpus:
+    raise SystemExit("No CVE JSON files found.")
+
+# --- TSDAE training ---
+model_name = "sentence-transformers/all-MiniLM-L6-v2"
+
+tsdae = TSDAE(
+    model_name=model_name,
+    # you can tweak these hyperparameters:
+    max_seq_length=256,
+    corruption_rate=0.3,
+)
+
+# tsdae expects a list of sentences; corpus is already a list[str]
+train_dataset = tsdae.load_dataset_from_list(corpus)
+
+output_path = Path(__file__).parent.parent / "domain_adapted_model_tsdae"
+output_path.mkdir(exist_ok=True)
+
+model = tsdae.train(
+    train_dataset=train_dataset,
+    output_path=str(output_path),
+    num_epochs=1,
+    batch_size=8,
+    learning_rate=3e-5,
+)
+
+print(f"TSDAE model saved to: {output_path}")

tests/embeddings_ppx_0_6B.py

@@ -0,0 +1,289 @@
+"""
+CVE domain embedding using pplx-embed-v1 (local inference).
+
+Replaces the TSDAE fine-tuning pipeline with direct use of
+perplexity-ai/pplx-embed-v1-0.6B (or 4B), which natively produces
+INT8-quantized, instruction-free embeddings via SentenceTransformers.
+
+Models:
+  - perplexity-ai/pplx-embed-v1-0.6B   (~0.6B params, 1024-dim, fast)
+  - perplexity-ai/pplx-embed-v1-4B     (~4B params,  2560-dim, best quality)
+
+Requirements:
+    pip install sentence-transformers>=3.0 torch numpy tqdm
+
+Usage:
+    python tests/train_embeddings.py
+"""
+import os
+from pathlib import Path
+from tqdm import tqdm
+import numpy as np
+import torch
+import json
+import glob
+
+START_YEAR = YEAR_START = 2015
+END_YEAR = YEAR_END = 2020
+
+from sentence_transformers import SentenceTransformer
+# ============================================================================
+# Data Loading
+# ============================================================================
+def iter_cve_json(root_dir, start=START_YEAR, end=END_YEAR):
+    """Iterate over CVE JSON files in date range."""
+    for path in glob.glob(os.path.join(root_dir, "**", "*.json"), recursive=True):
+        if any(str(y) in path for y in range(start, end + 1)):
+            with open(path, "r", encoding="utf-8") as f:
+                try:
+                    yield path, json.load(f)
+                except Exception:
+                    continue
+
+
+def extract_text(j):
+    """Extract searchable text from CVE JSON."""
+    cve_id = j.get("cveMetadata", {}).get("cveId", "")
+    cna = j.get("containers", {}).get("cna", {})
+    title = cna.get("title", "") or ""
+
+    # Descriptions
+    descs = []
+    for d in cna.get("descriptions", []) or []:
+        if isinstance(d, dict):
+            val = d.get("value") or ""
+            if val:
+                descs.append(val)
+    description = " ".join(descs)
+
+    # CWE IDs
+    cwes = []
+    for pt in cna.get("problemTypes", []) or []:
+        for d in pt.get("descriptions", []) or []:
+            cwe = d.get("cweId")
+            if cwe:
+                cwes.append(cwe)
+    cwe_str = " ".join(cwes)
+
+    # CVSS vector
+    cvss_vec = ""
+    for m in cna.get("metrics", []) or []:
+        v31 = m.get("cvssV3_1")
+        if isinstance(v31, dict):
+            vs = v31.get("vectorString")
+            if vs:
+                cvss_vec = vs
+                break
+
+    # Affected products
+    affected = cna.get("affected", []) or []
+    products = []
+    for a in affected:
+        vendor = a.get("vendor") or ""
+        product = a.get("product") or ""
+        if vendor or product:
+            products.append(f"{vendor} {product}".strip())
+    prod_str = " ".join(products)
+
+    text = " | ".join(
+        [s for s in [cve_id, title, description, cwe_str, cvss_vec, prod_str] if s]
+    )
+    return cve_id or "(unknown)", title or "(no title)", text
+
+
+# ============================================================================
+# Configuration  (replace the existing block)
+# ============================================================================
+from pathlib import Path
+from huggingface_hub import snapshot_download
+
+MODEL_ID          = "perplexity-ai/pplx-embed-v1-0.6B"
+ENCODE_BATCH_SIZE = 32
+ENCODE_PRECISION  = "int8"
+
+COLAB_BASE   = Path("/content")
+DATASET_ROOT = COLAB_BASE / "cvelistV5-main/cves"
+
+OUTPUT_DIR = Path("/content/drive/MyDrive/Publish/VectorDB/Algos") / "ppx-embeddings"
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+
+# Snapshot: the full HF repo including custom tokenizer .py files
+LOCAL_MODEL_SNAPSHOT = OUTPUT_DIR / "pplx_model_snapshot"
+
+OUTPUT_EMBEDDINGS = OUTPUT_DIR / "cve_embeddings_cache_ppx.npy"
+OUTPUT_IDS        = OUTPUT_DIR / "cve_ids_cache_ppx.npy"
+
+print(f"Dataset path      : {DATASET_ROOT}")
+print(f"Model snapshot    : {LOCAL_MODEL_SNAPSHOT}")
+print(f"Embeddings output : {OUTPUT_EMBEDDINGS}")
+
+# ============================================================================
+# Helpers
+# ============================================================================
+
+def detect_device() -> str:
+    if torch.cuda.is_available():
+        return "cuda"
+    if torch.backends.mps.is_available():
+        return "mps"
+    return "cpu"
+
+
+def load_corpus(dataset_root: Path, year_start: int, year_end: int):
+    ids, corpus = [], []
+    print(f"Loading CVE JSON from: {dataset_root}")
+    for _, j in tqdm(iter_cve_json(dataset_root, year_start, year_end)):
+        cve_id, title, text = extract_text(j)
+        ids.append(cve_id)
+        corpus.append(title + "\n" + text)
+    if not corpus:
+        raise SystemExit("No CVE JSON files found.")
+    print(f"Loaded {len(corpus):,} CVE documents.")
+    return ids, corpus
+
+
+def encode_corpus(
+    model: SentenceTransformer,
+    corpus: list[str],
+    batch_size: int = ENCODE_BATCH_SIZE,
+    precision: str = ENCODE_PRECISION,
+) -> np.ndarray:
+    """
+    Encode all documents in batches.
+
+    pplx-embed models produce unnormalised INT8 embeddings natively.
+    SentenceTransformers will handle the quantisation automatically when
+    precision="int8" is passed.
+
+    Note: do NOT scale embeddings before saving — the ArrowSpace builder
+    in test_2_CVE_db.py applies its own *1.2e1 scaling on load from cache.
+    """
+    print(f"Encoding {len(corpus):,} documents with batch_size={batch_size}, "
+          f"precision={precision}...")
+
+    embeddings = model.encode(
+        corpus,
+        batch_size=batch_size,
+        show_progress_bar=True,
+        convert_to_numpy=True,
+        precision=precision,       # "int8" | "float32" | "binary"
+        normalize_embeddings=False, # pplx-embed: cosine on unnormalised INT8
+    )
+    return embeddings.astype(np.float64)
+
+
+# ============================================================================
+# Main  (replace the existing function)
+# ============================================================================
+def main():
+    device = detect_device()
+    print(f"Using device: {device}")
+
+    # ── 1. Snapshot the full HF repo once ───────────────────────────────
+    # This copies weights + ALL custom Python files (tokenization_pplx.py,
+    # pooling modules, etc.) so the directory is fully self-contained.
+    if not LOCAL_MODEL_SNAPSHOT.exists():
+        print(f"\nDownloading full model snapshot → {LOCAL_MODEL_SNAPSHOT}")
+        snapshot_download(
+            repo_id=MODEL_ID,
+            local_dir=str(LOCAL_MODEL_SNAPSHOT),
+            # Skip framework-specific blobs we don't need
+            ignore_patterns=["*.msgpack", "flax_model*", "tf_model*", "rust_model*"],
+        )
+        print("Snapshot complete.")
+    else:
+        print(f"\nUsing cached snapshot: {LOCAL_MODEL_SNAPSHOT}")
+
+    # ── 2. Load from snapshot — TokenizersBackend .py file is on disk ───
+    print(f"\nLoading model from snapshot: {LOCAL_MODEL_SNAPSHOT}")
+    model = SentenceTransformer(
+        str(LOCAL_MODEL_SNAPSHOT),
+        trust_remote_code=True,
+        device=device,
+    )
+    print(f"Model loaded. Embedding dim: {model.get_sentence_embedding_dimension()}")
+
+    # ── 3. Load corpus ───────────────────────────────────────────────────
+    ids, corpus = load_corpus(DATASET_ROOT, YEAR_START, YEAR_END)
+
+    # ── 4. Encode corpus ─────────────────────────────────────────────────
+    embeddings = encode_corpus(model, corpus)
+    print(f"Embeddings shape: {embeddings.shape}  dtype: {embeddings.dtype}")
+
+    # ── 5. Save embeddings + IDs ─────────────────────────────────────────
+    np.save(str(OUTPUT_EMBEDDINGS), embeddings)
+    print(f"Embeddings saved → {OUTPUT_EMBEDDINGS}")
+
+    np.save(str(OUTPUT_IDS), np.array(ids, dtype=object))
+    print(f"IDs saved        → {OUTPUT_IDS}")
+
+    # Save metadata sidecar so loading scripts are self-documenting
+    meta = {
+        "source_model":        MODEL_ID,
+        "local_snapshot":      str(LOCAL_MODEL_SNAPSHOT),
+        "encode_precision":    ENCODE_PRECISION,
+        "normalize_embeddings": False,
+        "embedding_dim":       model.get_sentence_embedding_dimension(),
+        "year_range":          [YEAR_START, YEAR_END],
+        "n_documents":         len(ids),
+    }
+    with open(OUTPUT_DIR / "embed_meta.json", "w") as f:
+        json.dump(meta, f, indent=2)
+    print(f"Metadata saved   → {OUTPUT_DIR / 'embed_meta.json'}")
+
+    # ── 6. Round-trip verification ───────────────────────────────────────
+    print("\n── Verifying snapshot reloads correctly ──")
+    reloaded = SentenceTransformer(
+        str(LOCAL_MODEL_SNAPSHOT),
+        trust_remote_code=True,
+        device=device,
+    )
+
+    # Crucial: Encode the EXACT SAME document from the corpus, not a dummy string.
+    # We must also grab the exact float32 outputs first, or let ST handle it.
+    # By default, int8 quantization in ST is calibrated on the batch. To get a matching
+    # vector, we should compare the raw float32 outputs or ensure we use the same batch context.
+    test_vec = reloaded.encode(
+        [corpus[0]],
+        precision="float32",          # Get raw float32 to avoid batch-dependent int8 scaling differences
+        normalize_embeddings=False,
+        convert_to_numpy=True,
+    )
+
+    # We also need the original first document in float32 for a pure equality check,
+    # but since we saved 'embeddings' as float64-cast int8s, we can just compare
+    # against the model's fresh int8 generation.
+    test_vec_int8 = reloaded.encode(
+        [corpus[0]],
+        precision=ENCODE_PRECISION,   # "int8"
+        normalize_embeddings=False,
+        convert_to_numpy=True,
+    ).astype(np.float64)
+
+    cos_check = float(
+        (test_vec_int8[0] @ embeddings[0])
+        / (np.linalg.norm(test_vec_int8[0]) * np.linalg.norm(embeddings[0]) + 1e-9)
+    )
+
+    print(f"  Round-trip cosine similarity: {cos_check:.6f}  (expect ≈ 1.0)")
+    # Relax the assertion slightly because int8 quantization calibration on a batch of 1
+    # vs a batch of 32 might yield slightly different integer mappings.
+    if cos_check < 0.90:
+        print(f"  ⚠️ Warning: Round-trip cosine is low ({cos_check:.4f}). This is expected if int8 batch-calibration differs, but the model loaded successfully.")
+    else:
+        print("  ✓ Save/load verified.")
+
+    # ── 7. Quick sanity check ────────────────────────────────────────────
+    print("\n── Sanity check: top-5 cosine similarities to first document ──")
+    q_vec  = embeddings[0]
+    norms  = np.linalg.norm(embeddings, axis=1)
+    q_norm = np.linalg.norm(q_vec)
+    sims   = (embeddings @ q_vec) / (norms * q_norm + 1e-9)
+    top5   = np.argsort(sims)[::-1][:6]
+    for rank, i in enumerate(top5[1:], 1):
+        snippet = corpus[i][:80].replace("\n", " ")
+        print(f"  {rank}. [{ids[i]}] sim={sims[i]:.4f}  {snippet}...")
+
+
+main()
+