Wire modelatlas_similar to HF parquet dataset (2,435 models)

requirements.txt

@@ -1,2 +1,5 @@
 requests>=2.31.0
 psycopg2-binary>=2.9.9
+pandas>=2.0.0
+pyarrow>=14.0.0
+huggingface_hub>=0.23.0

scan.py

@@ -12,10 +12,51 @@ from datetime import datetime, timezone
 from pathlib import Path
 from typing import Optional
 import requests
-import psycopg2, psycopg2.extras
 
-DB = "postgresql:///modelatlas?host=/var/run/postgresql&port=5433&user=tim"
 HF_API = "https://huggingface.co"
+HF_DATASET = "RadicalNotionAI/modelatlas-reference"
+DB = "postgresql:///modelatlas?host=/var/run/postgresql&port=5433&user=tim"
+
+# In-process cache — loaded once per worker, refreshes when the file changes
+_REF_DF = None
+_REF_LOADED_AT: float = 0.0
+_REF_TTL = 3600  # reload at most once per hour
+
+
+def _load_reference_df():
+    """Load ModelAtlas reference parquet. Tries local snapshot first, then HF dataset."""
+    global _REF_DF, _REF_LOADED_AT
+    now = time.time()
+    if _REF_DF is not None and (now - _REF_LOADED_AT) < _REF_TTL:
+        return _REF_DF
+
+    import pandas as pd
+
+    # 1. Local snapshot (fast, used in dev / on local server)
+    local_path = Path(__file__).parent.parent / "snapshots" / "modeldna_reference.parquet"
+    if local_path.exists():
+        try:
+            _REF_DF = pd.read_parquet(local_path)
+            _REF_LOADED_AT = now
+            return _REF_DF
+        except Exception:
+            pass
+
+    # 2. HF dataset (used on HF Space — downloaded and cached by huggingface_hub)
+    try:
+        from huggingface_hub import hf_hub_download
+        path = hf_hub_download(
+            repo_id=HF_DATASET,
+            filename="modeldna_reference.parquet",
+            repo_type="dataset",
+        )
+        _REF_DF = pd.read_parquet(path)
+        _REF_LOADED_AT = now
+        return _REF_DF
+    except Exception:
+        pass
+
+    return None
 
 # Known base model reference configs (canonical identifiers)
 KNOWN_BASES = {
@@ -196,26 +237,44 @@ def stage1_screen(model_id: str, config: dict) -> dict:
                 "evidence": reasons,
             })
 
-    # Check ModelAtlas DB for exact signature
+    # Query ModelAtlas reference parquet for architecturally similar models
     db_matches = []
     try:
-        conn = psycopg2.connect(DB)
-        cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
-        cur.execute("""
-            SELECT m.model_id, o.name AS lab, m.hf_downloads, m.release_date,
-                   a.technique_signature, a.total_params, a.num_layers, a.hidden_size, a.vocab_size
-            FROM analyses a JOIN models m ON m.id=a.model_id
-            JOIN organizations o ON m.org_id=o.id
-            WHERE a.is_current=true AND a.vocab_size=%s AND a.hidden_size=%s
-              AND m.model_id NOT ILIKE '%%tiny%%' AND m.model_id NOT ILIKE '/%%'
-            ORDER BY m.hf_downloads DESC NULLS LAST
-            LIMIT 5
-        """, (vocab, hidden))
-        db_matches = [dict(r) for r in cur.fetchall()]
-        cur.close(); conn.close()
+        ref = _load_reference_df()
+        if ref is not None and vocab and hidden:
+            hit = ref[
+                (ref["vocab_size"] == vocab) &
+                (ref["hidden_size"] == hidden) &
+                (~ref["model_id"].str.contains("tiny|/", case=False, na=False))
+            ].sort_values("hf_downloads", ascending=False).head(5)
+            db_matches = hit[
+                ["model_id", "org_display", "hf_downloads", "total_params",
+                 "technique_signature", "num_layers", "hidden_size", "vocab_size"]
+            ].rename(columns={"org_display": "lab"}).to_dict("records")
     except Exception:
         pass
 
+    # Also try local DB if available (dev / local server)
+    if not db_matches:
+        try:
+            import psycopg2, psycopg2.extras
+            conn = psycopg2.connect(DB)
+            cur = conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
+            cur.execute("""
+                SELECT m.model_id, o.name AS lab, m.hf_downloads, m.release_date,
+                       a.technique_signature, a.total_params, a.num_layers, a.hidden_size, a.vocab_size
+                FROM analyses a JOIN models m ON m.id=a.model_id
+                JOIN organizations o ON m.org_id=o.id
+                WHERE a.is_current=true AND a.vocab_size=%s AND a.hidden_size=%s
+                  AND m.model_id NOT ILIKE '%%tiny%%' AND m.model_id NOT ILIKE '/%%'
+                ORDER BY m.hf_downloads DESC NULLS LAST
+                LIMIT 5
+            """, (vocab, hidden))
+            db_matches = [dict(r) for r in cur.fetchall()]
+            cur.close(); conn.close()
+        except Exception:
+            pass
+
     return {
         "arch_signature": arch_sig,
         "config_signals": {