254 lines
9.1 KiB
Python
254 lines
9.1 KiB
Python
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",
|
|
}
|