Blame for scripts/preflight_colbert.py

"""Build-time smoke check of the actual pinned PyLate mLateOn checkpoint.

Fail the Docker build before production if load, both encoding modes,
or 128D output are incompatible with the pinned package versions.
"""
from app.models import ColbertEngine


def main():
    engine = ColbertEngine.load()
    for kind in ("query", "passage"):
        result = engine.encode("Hello שלום привет", kind)
        assert result.token_embeddings, f"Empty mLateOn {kind} result"
        assert all(len(row) == 128 for row in result.token_embeddings), (
            f"Incorrect mLateOn {kind} vector shape"
        )
        print(f"mLateOn preflight {kind} passed: {result.tokens_count} tokens", flush=True)


if __name__ == "__main__":
    main()