Upload scripts/train_poster_sentry.py with huggingface_hub

scripts/train_poster_sentry.py

@@ -0,0 +1,402 @@
+#!/usr/bin/env python3
+"""
+Train PosterSentry on the real posters.science corpus.
+
+Data sources (all real, zero synthetic):
+    Positive (poster):
+        28K+ verified scientific posters from Zenodo & Figshare
+        /home/joneill/Nextcloud/vaults/jmind/calmi2/poster_science/poster-pdf-meta/downloads/
+
+    Negative (non_poster):
+        2,036 verified non-posters (multi-page docs, proceedings, abstracts)
+        Listed in: poster_classifier/non_posters_20251208_152217.txt
+
+        Plus: single pages extracted from armanc/scientific_papers (real papers)
+        Plus: ag_news articles (real junk text, rendered to match)
+
+Usage:
+    cd /home/joneill/pubverse_brett/poster_sentry
+    pip install -e ".[train]"
+    python scripts/train_poster_sentry.py --n-per-class 5000
+"""
+
+import argparse
+import json
+import logging
+import os
+import random
+import sys
+import time
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s | %(levelname)s | %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger = logging.getLogger(__name__)
+
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+
+# ── Paths ────────────────────────────────────────────────────────
+
+POSTER_SCIENCE_BASE = Path(
+    "/home/joneill/Nextcloud/vaults/jmind/calmi2/poster_science"
+)
+DOWNLOADS_DIR = POSTER_SCIENCE_BASE / "poster-pdf-meta" / "downloads"
+NON_POSTERS_LIST = (
+    POSTER_SCIENCE_BASE
+    / "poster_classifier"
+    / "non_posters_20251208_152217.txt"
+)
+CLASSIFICATION_JSON = (
+    POSTER_SCIENCE_BASE
+    / "poster_classifier"
+    / "classification_results_20251208_152217.json"
+)
+
+
+def _fix_path(p: str) -> str:
+    """Fix paths from classification JSON — they use /home/joneill/vaults/
+    but the actual Nextcloud mount is /home/joneill/Nextcloud/vaults/."""
+    if "/joneill/vaults/" in p and "/Nextcloud/" not in p:
+        return p.replace("/joneill/vaults/", "/joneill/Nextcloud/vaults/")
+    return p
+
+
+def collect_poster_paths(max_n: int = 10000) -> List[str]:
+    """Collect verified poster PDF paths from the corpus."""
+    # Load the classification results to get confirmed poster paths
+    if CLASSIFICATION_JSON.exists():
+        logger.info(f"Loading classification results from {CLASSIFICATION_JSON}")
+        with open(CLASSIFICATION_JSON) as f:
+            data = json.load(f)
+        poster_entries = data.get("posters", [])
+        paths = [_fix_path(e["pdf_path"]) for e in poster_entries if Path(_fix_path(e["pdf_path"])).exists()]
+        logger.info(f"  Found {len(paths)} verified poster paths")
+    else:
+        # Fallback: glob the downloads directory
+        logger.info(f"Globbing {DOWNLOADS_DIR} for PDFs...")
+        paths = [str(p) for p in DOWNLOADS_DIR.rglob("*.pdf")]
+        paths += [str(p) for p in DOWNLOADS_DIR.rglob("*.PDF")]
+        logger.info(f"  Found {len(paths)} PDFs")
+
+    random.shuffle(paths)
+    return paths[:max_n]
+
+
+def collect_non_poster_paths(max_n: int = 2000) -> List[str]:
+    """Collect verified non-poster PDF paths.
+
+    The non-posters were separated into:
+        poster-pdf-meta/separated_non_posters/downloads/{zenodo,figshare}/
+    """
+    paths = []
+
+    # Primary: glob the separated_non_posters directory
+    sep_dir = POSTER_SCIENCE_BASE / "poster-pdf-meta" / "separated_non_posters" / "downloads"
+    if sep_dir.exists():
+        for pdf in sep_dir.rglob("*.pdf"):
+            paths.append(str(pdf))
+        for pdf in sep_dir.rglob("*.PDF"):
+            paths.append(str(pdf))
+        logger.info(f"  Found {len(paths)} non-poster PDFs in {sep_dir}")
+    else:
+        # Fallback: try the original list with path fixing
+        logger.info("  Separated dir not found, trying original list...")
+        if NON_POSTERS_LIST.exists():
+            with open(NON_POSTERS_LIST) as f:
+                for line in f:
+                    p = _fix_path(line.strip())
+                    if p and Path(p).exists():
+                        paths.append(p)
+            logger.info(f"  Found {len(paths)} verified non-poster paths from list")
+
+    random.shuffle(paths)
+    return paths[:max_n]
+
+
+def extract_features_from_pdfs(
+    pdf_paths: List[str],
+    label: int,
+    text_model,
+    visual_ext,
+    structural_ext,
+    max_text_chars: int = 4000,
+) -> Tuple[np.ndarray, np.ndarray, List[str]]:
+    """
+    Extract multimodal features from a list of PDFs.
+
+    Returns (X, y, extracted_texts) where:
+        X: (N, 542) feature matrix
+        y: (N,) labels
+        extracted_texts: list of extracted text strings (for PubGuard reuse)
+    """
+    from tqdm import tqdm
+    import fitz
+    import re
+
+    embeddings = []
+    visual_vecs = []
+    struct_vecs = []
+    texts_out = []
+    labels = []
+
+    for pdf_path in tqdm(pdf_paths, desc=f"{'poster' if label == 1 else 'non_poster'}"):
+        try:
+            # Extract text
+            doc = fitz.open(pdf_path)
+            if len(doc) == 0:
+                doc.close()
+                continue
+            text = doc[0].get_text()
+            doc.close()
+            text = re.sub(r"\s+", " ", text).strip()[:max_text_chars]
+
+            if len(text) < 20:
+                continue
+
+            # Visual features
+            img = visual_ext.pdf_to_image(pdf_path)
+            if img is not None:
+                vf = visual_ext.extract(img)
+            else:
+                vf = {n: 0.0 for n in visual_ext.FEATURE_NAMES}
+
+            # Structural features
+            sf = structural_ext.extract(pdf_path)
+
+            texts_out.append(text)
+            visual_vecs.append(visual_ext.to_vector(vf))
+            struct_vecs.append(structural_ext.to_vector(sf))
+            labels.append(label)
+
+        except Exception as e:
+            logger.debug(f"Skipping {pdf_path}: {e}")
+            continue
+
+    if not texts_out:
+        return np.array([]), np.array([]), []
+
+    # Embed all texts at once
+    logger.info(f"Embedding {len(texts_out)} texts...")
+    emb = text_model.encode(texts_out, show_progress_bar=True)
+    norms = np.linalg.norm(emb, axis=1, keepdims=True)
+    norms = np.where(norms == 0, 1, norms)
+    emb = (emb / norms).astype("float32")
+
+    visual_arr = np.array(visual_vecs, dtype="float32")
+    struct_arr = np.array(struct_vecs, dtype="float32")
+
+    X = np.concatenate([emb, visual_arr, struct_arr], axis=1)
+    y = np.array(labels)
+
+    return X, y, texts_out
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Train PosterSentry")
+    parser.add_argument("--n-per-class", type=int, default=5000,
+                        help="Max samples per class (poster/non_poster)")
+    parser.add_argument("--test-size", type=float, default=0.15)
+    parser.add_argument("--models-dir", default=None)
+    parser.add_argument("--export-texts", default=None,
+                        help="Export extracted texts as NDJSON for PubGuard retraining")
+    args = parser.parse_args()
+
+    from model2vec import StaticModel
+    from sklearn.linear_model import LogisticRegression
+    from sklearn.metrics import classification_report
+    from sklearn.model_selection import train_test_split
+    from poster_sentry.features import VisualFeatureExtractor, PDFStructuralExtractor
+
+    # Models dir
+    if args.models_dir:
+        models_dir = Path(args.models_dir)
+    else:
+        models_dir = Path.home() / ".poster_sentry" / "models"
+    models_dir.mkdir(parents=True, exist_ok=True)
+
+    # Load embedding model
+    logger.info("Loading model2vec...")
+    emb_cache = models_dir / "poster-sentry-embedding"
+    if emb_cache.exists():
+        text_model = StaticModel.from_pretrained(str(emb_cache))
+    else:
+        text_model = StaticModel.from_pretrained("minishlab/potion-base-32M")
+        emb_cache.parent.mkdir(parents=True, exist_ok=True)
+        text_model.save_pretrained(str(emb_cache))
+
+    visual_ext = VisualFeatureExtractor()
+    structural_ext = PDFStructuralExtractor()
+
+    # ── Collect data ─────────────────────────────────────────────
+    logger.info("=" * 60)
+    logger.info("Collecting training data...")
+    logger.info("=" * 60)
+
+    poster_paths = collect_poster_paths(max_n=args.n_per_class)
+    non_poster_paths = collect_non_poster_paths(max_n=args.n_per_class)
+
+    logger.info(f"Poster PDFs to process: {len(poster_paths)}")
+    logger.info(f"Non-poster PDFs to process: {len(non_poster_paths)}")
+
+    # ── Extract features ─────────────────────────────────────────
+    logger.info("=" * 60)
+    logger.info("Extracting features from poster PDFs...")
+    logger.info("=" * 60)
+
+    X_pos, y_pos, texts_pos = extract_features_from_pdfs(
+        poster_paths, label=1, text_model=text_model,
+        visual_ext=visual_ext, structural_ext=structural_ext,
+    )
+
+    logger.info(f"Poster features: {X_pos.shape}")
+
+    logger.info("=" * 60)
+    logger.info("Extracting features from non-poster PDFs...")
+    logger.info("=" * 60)
+
+    X_neg, y_neg, texts_neg = extract_features_from_pdfs(
+        non_poster_paths, label=0, text_model=text_model,
+        visual_ext=visual_ext, structural_ext=structural_ext,
+    )
+
+    logger.info(f"Non-poster features: {X_neg.shape}")
+
+    # ── Balance classes ──────────────────────────────────────────
+    min_count = min(len(y_pos), len(y_neg))
+    logger.info(f"Balancing: {min_count} samples per class")
+
+    if len(y_pos) > min_count:
+        idx = np.random.choice(len(y_pos), min_count, replace=False)
+        X_pos = X_pos[idx]
+        y_pos = y_pos[idx]
+        texts_pos = [texts_pos[i] for i in idx]
+
+    if len(y_neg) > min_count:
+        idx = np.random.choice(len(y_neg), min_count, replace=False)
+        X_neg = X_neg[idx]
+        y_neg = y_neg[idx]
+        texts_neg = [texts_neg[i] for i in idx]
+
+    X = np.vstack([X_pos, X_neg])
+    y = np.concatenate([y_pos, y_neg])
+
+    logger.info(f"Total training data: {X.shape} (poster={sum(y)}, non_poster={len(y)-sum(y)})")
+
+    # ── Export texts for PubGuard ────────────────────────────────
+    if args.export_texts:
+        export_path = Path(args.export_texts)
+        export_path.parent.mkdir(parents=True, exist_ok=True)
+        with open(export_path, "w") as f:
+            for text in texts_pos:
+                f.write(json.dumps({"text": text, "label": "poster"}) + "\n")
+            for text in texts_neg:
+                f.write(json.dumps({"text": text, "label": "non_poster"}) + "\n")
+        logger.info(f"Exported {len(texts_pos) + len(texts_neg)} texts to {export_path}")
+
+    # ── Feature scaling ──────────────────────────────────────────
+    # Critical: the 512-d text embedding drowns out the 30 structural/visual
+    # features if we don't scale. StandardScaler normalizes each column to
+    # zero mean and unit variance, giving structural signals fair weight.
+    from sklearn.preprocessing import StandardScaler
+
+    logger.info("=" * 60)
+    logger.info("Scaling features (StandardScaler)")
+    logger.info("=" * 60)
+
+    scaler = StandardScaler()
+    X_scaled = scaler.fit_transform(X)
+
+    # Log feature variance to confirm structural features are alive
+    emb_var = np.mean(np.var(X_scaled[:, :512], axis=0))
+    vis_var = np.mean(np.var(X_scaled[:, 512:527], axis=0))
+    str_var = np.mean(np.var(X_scaled[:, 527:], axis=0))
+    logger.info(f"  Mean variance — text: {emb_var:.3f}  visual: {vis_var:.3f}  structural: {str_var:.3f}")
+
+    # ── Train ────────────────────────────────────────────────────
+    logger.info("=" * 60)
+    logger.info("Training PosterSentry classifier")
+    logger.info("=" * 60)
+
+    X_tr, X_te, y_tr, y_te = train_test_split(
+        X_scaled, y, test_size=args.test_size, stratify=y, random_state=SEED,
+    )
+
+    logger.info(f"Train: {X_tr.shape[0]:,} | Test: {X_te.shape[0]:,}")
+    logger.info(f"Features: {X_tr.shape[1]} (512 text + 15 visual + 15 structural)")
+
+    clf = LogisticRegression(
+        C=1.0, max_iter=1000, class_weight="balanced",
+        solver="lbfgs", n_jobs=1, random_state=SEED,
+    )
+
+    t0 = time.time()
+    clf.fit(X_tr, y_tr)
+    elapsed = time.time() - t0
+    logger.info(f"Trained in {elapsed:.1f}s")
+
+    y_pred = clf.predict(X_te)
+    labels = ["non_poster", "poster"]
+    report = classification_report(y_te, y_pred, target_names=labels, digits=4)
+    logger.info(f"\n{report}")
+
+    # Show top feature importances
+    coef = clf.coef_[0]
+    all_names = (
+        [f"emb_{i}" for i in range(512)]
+        + list(VisualFeatureExtractor.FEATURE_NAMES)
+        + list(PDFStructuralExtractor.FEATURE_NAMES)
+    )
+    top_idx = np.argsort(np.abs(coef))[-15:][::-1]
+    logger.info("Top 15 features by |coefficient|:")
+    for idx in top_idx:
+        logger.info(f"  {all_names[idx]:30s}  coef={coef[idx]:+.4f}")
+
+    # ── Save head as .npz ────────────────────────────────────────
+    if clf.coef_.shape[0] == 1:
+        W = np.vstack([-clf.coef_[0], clf.coef_[0]]).T.astype("float32")
+        b = np.array([-clf.intercept_[0], clf.intercept_[0]], dtype="float32")
+    else:
+        W = clf.coef_.T.astype("float32")
+        b = clf.intercept_.astype("float32")
+
+    head_path = models_dir / "poster_sentry_head.npz"
+    np.savez(
+        head_path, W=W, b=b, labels=np.array(labels),
+        scaler_mean=scaler.mean_.astype("float32"),
+        scaler_scale=scaler.scale_.astype("float32"),
+    )
+    logger.info(f"Saved classifier head + scaler → {head_path}")
+
+    # ── Smoke test ───────────────────────────────────────────────
+    logger.info("\n" + "=" * 60)
+    logger.info("SMOKE TEST")
+    logger.info("=" * 60)
+
+    from poster_sentry import PosterSentry
+
+    sentry = PosterSentry(models_dir=models_dir)
+    sentry.initialize()
+
+    # Test with some real PDFs
+    test_pdfs = poster_paths[:2] + non_poster_paths[:2]
+    for p in test_pdfs:
+        try:
+            result = sentry.classify(p)
+            icon = "📋" if result["is_poster"] else "📄"
+            print(f"  {icon} {Path(p).name[:60]:60s}  poster={result['is_poster']}  conf={result['confidence']:.3f}")
+        except Exception as e:
+            print(f"  ⚠️  {Path(p).name[:60]}: {e}")
+
+    logger.info(f"\nDone! Model saved to: {models_dir}")
+
+
+if __name__ == "__main__":
+    main()