Cleanup init commit

core/registry.py

@@ -0,0 +1,253 @@
+import json
+import re
+import time
+
+from huggingface_hub import HfApi
+
+import config
+
+HUB_SEED_REPO = "neuronpedia/jacobian-lens"
+HUB_CACHE_TTL = 600
+
+_hub_cache = {"at": 0.0, "entries": None, "error": None}
+
+
+def local_lenses():
+    out = []
+    if not config.LENSES_DIR.exists():
+        return out
+    for entry in sorted(config.LENSES_DIR.iterdir()):
+        lens_file = entry / "lens.pt"
+        if not entry.is_dir() or not lens_file.exists():
+            continue
+        meta = {}
+        meta_file = entry / "meta.json"
+        if meta_file.exists():
+            meta = json.loads(meta_file.read_text(encoding="utf-8"))
+        out.append({"name": entry.name, "path": str(lens_file), "meta": meta})
+    return out
+
+
+def _base_model_from(filename):
+    stem = filename.rsplit("/", 1)[-1].removesuffix(".pt")
+    match = re.match(r"(.+?)_jacobian_lens(?:_n\d+)?$", stem)
+    return match.group(1) if match else None
+
+
+def _derived_model_id(base):
+    if base is None:
+        return None
+    lowered = base.lower()
+    if lowered.startswith("qwen"):
+        return f"Qwen/{base}"
+    if lowered.startswith("gemma"):
+        return f"google/{base}"
+    if lowered.startswith("llama"):
+        return f"meta-llama/{base}"
+    if lowered.startswith("gpt-oss"):
+        return f"openai/{base}"
+    if lowered == "gpt2":
+        return "openai-community/gpt2"
+    if lowered.startswith("pythia"):
+        return f"EleutherAI/{base}"
+    if lowered.startswith("olmo"):
+        return f"allenai/{base}"
+    return base
+
+
+def hub_lenses(force=False):
+    now = time.time()
+    if not force and _hub_cache["entries"] is not None and now - _hub_cache["at"] < HUB_CACHE_TTL:
+        return _hub_cache["entries"]
+    api = HfApi()
+    repos = {HUB_SEED_REPO}
+    try:
+        for model in api.list_models(search="jacobian-lens", limit=50):
+            repos.add(model.id)
+        for model in api.list_models(filter="jacobian_lens", limit=50):
+            repos.add(model.id)
+    except Exception as exc:
+        _hub_cache.update(error=f"Hub search unavailable: {exc}")
+    entries = []
+    for repo_id in sorted(repos):
+        try:
+            refs = api.list_repo_refs(repo_id)
+            branches = [b.name for b in refs.branches] or ["main"]
+        except Exception:
+            continue
+        for branch in branches:
+            try:
+                files = api.list_repo_files(repo_id, revision=branch)
+            except Exception:
+                continue
+            for filename in files:
+                if not filename.endswith(".pt"):
+                    continue
+                base = _base_model_from(filename)
+                entries.append(
+                    {
+                        "repo_id": repo_id,
+                        "revision": branch,
+                        "filename": filename,
+                        "base_model": base,
+                        "derived_model_id": _derived_model_id(base),
+                        "model_revision_verified": False,
+                    }
+                )
+    _hub_cache.update(at=now, entries=entries)
+    return entries
+
+
+_base_cache = {}
+BASE_CACHE_TTL = 600
+
+
+def hub_base_model(model_id):
+    """Base model declared by the repo's model card (``base_model`` tags) —
+    e.g. a finetune pointing at the checkpoint it was trained from. ``None``
+    if unknown, offline, or not a Hub repo."""
+    now = time.time()
+    hit = _base_cache.get(model_id)
+    if hit and now - hit["at"] < BASE_CACHE_TTL:
+        return hit["base"]
+    found = None
+    if "/" in model_id and not model_id.startswith("local/"):
+        try:
+            info = HfApi().model_info(model_id)
+            for tag in info.tags or []:
+                if not tag.startswith("base_model:"):
+                    continue
+                rest = tag[len("base_model:"):]
+                if ":" in rest:  # qualified form: finetune:X, adapter:X, quantized:X
+                    rest = rest.split(":", 1)[1]
+                if rest and rest.lower() != model_id.lower():
+                    found = rest
+                    break
+        except Exception:
+            found = None
+    _base_cache[model_id] = {"at": now, "base": found}
+    return found
+
+
+def lenses_for_model(model_id, revision=None):
+    matches_local = []
+    for lens in local_lenses():
+        meta = lens["meta"]
+        if meta.get("model_id") != model_id:
+            continue
+        lens_rev = meta.get("model_revision")
+        compatible = True
+        reason = None
+        if revision and lens_rev and lens_rev != revision:
+            compatible = False
+            reason = f"fit revision ({lens_rev[:12]}) != loaded model ({revision[:12]})"
+        elif lens_rev is None and not model_id.startswith("local/"):
+            reason = "fit revision unknown"
+        matches_local.append(dict(lens, compatible=compatible, reason=reason))
+
+    base = model_id.split("/")[-1].lower()
+    base_ref = hub_base_model(model_id)  # e.g. "google/gemma-3-1b-it" for a finetune
+    base_ref_name = base_ref.split("/")[-1].lower() if base_ref else None
+
+    def hub_entry(entry, via, reason=None, compatible=True):
+        return dict(
+            entry,
+            via=via,
+            compatible=compatible,
+            reason=reason,
+            cached=_lens_cached(entry["repo_id"], entry["filename"], entry["revision"]),
+        )
+
+    # one entry per branch in hub_lenses → dedupe, main first
+    entries = []
+    seen = set()
+    for entry in sorted(hub_lenses(), key=lambda e: e["revision"] != "main"):
+        if entry["base_model"] is None or (entry["repo_id"], entry["filename"]) in seen:
+            continue
+        seen.add((entry["repo_id"], entry["filename"]))
+        entries.append(entry)
+
+    matches_hub = []
+    matched = set()
+    prefix_hits = {}
+    for entry in entries:
+        key = (entry["repo_id"], entry["filename"])
+        name = entry["base_model"].lower()
+        derived = (entry["derived_model_id"] or "").lower()
+        if name == base or derived == model_id.lower():
+            # ⚠ only for a real problem (local merge); the fit revision not being
+            # published is the normal state of Hub repos → discreet note
+            reason = None
+            if model_id.startswith("local/"):
+                reason = "local model: a Hub lens fitted on the original checkpoint doesn't match a merge"
+            matched.add(key)
+            matches_hub.append(hub_entry(
+                entry, "model", reason, compatible=not model_id.startswith("local/")))
+        elif base_ref and (name == base_ref_name or derived == base_ref.lower()):
+            matched.add(key)
+            matches_hub.append(hub_entry(
+                entry, "base-model",
+                f"lens of the base model {base_ref} — fitted on the original "
+                "weights, a finetune's readouts may drift slightly"))
+        elif base_ref is None and name != base and len(name) >= 6 and base.startswith(name):
+            prefix_hits.setdefault(len(name), []).append(entry)
+
+    # No card metadata: fall back to the longest name prefix (a finetune usually
+    # keeps its base's name — "gemma-3-1b-it-toxicity" → "gemma-3-1b-it").
+    if prefix_hits and not any(m["via"] == "model" for m in matches_hub):
+        for entry in prefix_hits[max(prefix_hits)]:
+            matched.add((entry["repo_id"], entry["filename"]))
+            matches_hub.append(hub_entry(
+                entry, "base-guess",
+                f"the name suggests a finetune of {entry['base_model']} — fitted "
+                "on the original weights, readouts may drift slightly"))
+
+    # Everything else stays reachable for cross-model loading (your own
+    # architecture-compatible lens); d_model/layers are checked at load time.
+    others = [
+        hub_entry(entry, "other")
+        for entry in entries
+        if (entry["repo_id"], entry["filename"]) not in matched
+    ]
+    return {
+        "local": matches_local,
+        "hub": matches_hub,
+        "other": others,
+        "base_model": base_ref,
+        "hub_error": _hub_cache.get("error"),
+    }
+
+
+def _lens_cached(repo_id, filename, revision=None):
+    """True if the lens file is already in the local HF cache (no download on load)."""
+    from huggingface_hub import try_to_load_from_cache
+
+    try:
+        result = try_to_load_from_cache(repo_id, filename, revision=revision)
+        return isinstance(result, str)
+    except Exception:
+        return False
+
+
+def resolve_lens(path=None, repo_id=None, filename=None):
+    if path:
+        for lens in local_lenses():
+            if lens["path"] == path:
+                meta = lens["meta"]
+                return {
+                    "source": "local",
+                    "name": lens["name"],
+                    "required_model": meta.get("model_id"),
+                    "required_revision": meta.get("model_revision"),
+                    "meta": meta,
+                }
+        return {"source": "local", "required_model": None, "meta": {}, "warning": "meta.json missing: required model unknown"}
+    base = _base_model_from(filename or "")
+    return {
+        "source": "hub",
+        "repo_id": repo_id,
+        "filename": filename,
+        "required_model": _derived_model_id(base),
+        "required_revision": None,
+        "warning": "model derived from the filename; exact revision not published",
+    }