chore: script to fix PsyEmbedding HF repos' missing 1_Pooling config

Uploads the 1_Pooling/config.json that modules.json references but the four
Culture-and-Morality-Lab/psyembedding-* repos never shipped (the reason for
the pooling_fallback workaround). Dry-run by default; --write needs lab HF
credentials; --verify asserts auto-load == fallback assembly.

scripts/fix_psyembedding_pooling.py

@@ -0,0 +1,160 @@
+#!/usr/bin/env python3
+"""Fix the PsyEmbedding repos' broken sentence-transformers packaging.
+
+All four Culture-and-Morality-Lab/psyembedding-* repos ship a modules.json
+that references a 1_Pooling module, but the 1_Pooling/config.json it points
+at was never uploaded - so SentenceTransformer(repo_id) cannot auto-load
+them (the reason for the pooling_fallback workaround in models.yaml and
+backend/app/ccr.py). This uploads the missing config, which makes the repos
+load the standard way for everyone, platform or not.
+
+The pooling config is derived, not assumed: pooling mode comes from the
+registry entry (mean, per the model cards - the same thing the platform's
+fallback has been computing all along) and the embedding dimension is read
+from each repo's own config.json at the pinned revision. Adding the file
+changes NOTHING numerically: same weights, same pooling math, byte-identical
+embeddings - it only lets the packaged loader do what our fallback did.
+
+Usage (write needs lab credentials: `hf auth login` or HF_TOKEN env var):
+    python scripts/fix_psyembedding_pooling.py            # dry run: show plan
+    python scripts/fix_psyembedding_pooling.py --write    # commit to the repos
+    python scripts/fix_psyembedding_pooling.py --verify REPO_ID [REVISION]
+        # download one fixed repo and check auto-load == manual assembly
+
+After --write, finish in the platform repo (do NOT blanket-run
+pin_revisions.py --write; it would repin every model, not just these):
+    1. models.yaml: set the four psyembedding revision lines to the new SHAs
+       printed below, and delete their pooling_fallback lines.
+    2. python packages/model_registry/validate_models.py
+    3. cd backend && python -m pytest -q
+"""
+
+import json
+import sys
+
+from huggingface_hub import HfApi, hf_hub_download
+
+REPOS = [
+    "Culture-and-Morality-Lab/psyembedding-bert-large-uncased",
+    "Culture-and-Morality-Lab/psyembedding-roberta-large",
+    "Culture-and-Morality-Lab/psyembedding-gte-large",
+    "Culture-and-Morality-Lab/psyembedding-e5-large-v2",
+]
+
+COMMIT_MESSAGE = (
+    "Add missing 1_Pooling/config.json referenced by modules.json\n\n"
+    "modules.json declares a Pooling module at 1_Pooling, but the config "
+    "was never uploaded, so SentenceTransformer auto-loading fails. Mean "
+    "pooling per the model card; no change to weights or outputs."
+)
+
+# The classic 5-key Pooling config: loads on every sentence-transformers
+# version (newer optional flags default to the same behavior when absent).
+def pooling_config(dim: int) -> dict:
+    return {
+        "word_embedding_dimension": dim,
+        "pooling_mode_cls_token": False,
+        "pooling_mode_mean_tokens": True,
+        "pooling_mode_max_tokens": False,
+        "pooling_mode_mean_sqrt_len_tokens": False,
+    }
+
+
+def plan(api: HfApi, repo: str) -> dict | None:
+    """Validate assumptions against the live repo; return the upload plan."""
+    files = {f.path for f in api.list_repo_tree(repo)}
+    if "1_Pooling/config.json" in files:
+        print(f"  already fixed - skipping")
+        return None
+
+    modules = json.load(open(hf_hub_download(repo, "modules.json")))
+    pooling = [m for m in modules if m["type"].endswith("models.Pooling")]
+    assert pooling and pooling[0]["path"] == "1_Pooling", (
+        f"{repo}: modules.json does not reference 1_Pooling as expected: {modules}"
+    )
+
+    hf_config = json.load(open(hf_hub_download(repo, "config.json")))
+    dim = hf_config["hidden_size"]
+    assert dim == 1024, f"{repo}: hidden_size {dim} != 1024 in the registry"
+
+    content = json.dumps(pooling_config(dim), indent=2) + "\n"
+    print(f"  will add 1_Pooling/config.json (mean pooling, dim {dim})")
+    return {"repo": repo, "content": content}
+
+
+def write(api: HfApi, p: dict) -> None:
+    info = api.create_commit(
+        repo_id=p["repo"],
+        operations=[
+            __import__("huggingface_hub").CommitOperationAdd(
+                path_in_repo="1_Pooling/config.json",
+                path_or_fileobj=p["content"].encode(),
+            )
+        ],
+        commit_message=COMMIT_MESSAGE,
+    )
+    print(f"  committed: {info.oid}")
+
+
+def verify(repo: str, revision: str | None) -> None:
+    """Auto-load must equal the manual Transformer+Pooling assembly."""
+    import numpy as np
+    from sentence_transformers import SentenceTransformer
+    from sentence_transformers import models as st_models
+
+    texts = [
+        "I am deeply satisfied with my life.",
+        "The bus was late again this morning.",
+        "Caring for the vulnerable is the most important virtue.",
+    ]
+    auto = SentenceTransformer(repo, revision=revision)
+    word = st_models.Transformer(repo, max_seq_length=512)
+    get_dim = getattr(word, "get_embedding_dimension", None) or word.get_word_embedding_dimension
+    manual = SentenceTransformer(modules=[word, st_models.Pooling(get_dim(), pooling_mode="mean")])
+
+    a = auto.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
+    m = manual.encode(texts, convert_to_numpy=True, normalize_embeddings=True)
+    assert a.shape == m.shape, (a.shape, m.shape)
+    max_diff = float(np.abs(a - m).max())
+    print(f"{repo}: shapes {a.shape}, max |auto - manual| = {max_diff:.2e}")
+    assert max_diff == 0.0, "auto-load does not match the platform's fallback assembly"
+    print("verified: byte-identical to the fallback the platform has been using")
+
+
+def main() -> None:
+    if "--verify" in sys.argv:
+        i = sys.argv.index("--verify")
+        repo = sys.argv[i + 1]
+        rev = sys.argv[i + 2] if len(sys.argv) > i + 2 else None
+        verify(repo, rev)
+        return
+
+    do_write = "--write" in sys.argv
+    api = HfApi()
+    plans = []
+    for repo in REPOS:
+        print(repo)
+        p = plan(api, repo)
+        if p:
+            plans.append(p)
+
+    if not plans:
+        print("\nNothing to do.")
+        return
+    if not do_write:
+        print(f"\nDry run: {len(plans)} repo(s) need the fix. "
+              "Re-run with --write (requires lab HF credentials).")
+        print("Config that would be uploaded:\n" + plans[0]["content"])
+        return
+
+    who = api.whoami()
+    print(f"\nWriting as: {who['name']} ({who.get('email', '?')})")
+    for p in plans:
+        print(p["repo"])
+        write(api, p)
+    print("\nDone. New commit SHAs are above - now update models.yaml "
+          "(revision lines + remove pooling_fallback) per the module docstring.")
+
+
+if __name__ == "__main__":
+    main()