scripts/fix_psyembedding_pooling.py

#!/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()