"""Registry of embedding models for the SQ1 model screen.

Each entry: HF id, max sequence length the model supports, the prefix (or prompt) it expects on passages
and on queries, and the folder its shards live in under data/.  Shared by the Colab notebook, build_index.py
and search.py so every model is embedded and queried the way its authors intend.

    from embed_models import MODELS
    MODELS["bge"]  -> {"hf": "BAAI/bge-m3", "max_seq": 1024, ...}
"""
MODELS = {
    # key      HF model id                                    max_seq  passage prefix   query prefix                                   index dir     notes
    "bge":     {"hf": "BAAI/bge-m3",                              "max_seq": 1024, "passage": "",            "query": "",
                "dir": "index",      "note": "8k context; 100+ languages; current baseline"},
    "sealion": {"hf": "aisingapore/SEA-LION-E5-Embedding-600M",   "max_seq": 512,  "passage": "passage: ",  "query": "query: ",
                "dir": "index_sealion", "note": "mE5-large fine-tuned on 11 SEA languages incl. Burmese; 512-token limit"},
    "me5i":    {"hf": "intfloat/multilingual-e5-large-instruct",  "max_seq": 512,  "passage": "",
                "query": "Instruct: Given a question about Myanmar parliamentary debates, retrieve the passages that answer it\nQuery: ",
                "dir": "index_me5i", "note": "best Burmese score on SEA-BED; instruction on the query side only; 512-token limit"},
    "me5":     {"hf": "intfloat/multilingual-e5-large",           "max_seq": 512,  "passage": "passage: ",  "query": "query: ",
                "dir": "index_me5",  "note": "plain mE5; 512-token limit"},
    "modernbert": {"hf": "aisingapore/SEA-LION-ModernBERT-Embedding-600M", "max_seq": 1024, "passage": "", "query": "",
                "dir": "index_modernbert", "note": "8k context, SEA languages incl. Burmese; new, not yet on the leaderboard"},
    "myx":     {"hf": "DatarrX/myX-Semantic",                     "max_seq": 512,  "passage": "",            "query": "",
                "dir": "index_myx",  "note": "Burmese-only MiniLM distillation, 768-d; no published benchmarks; cheap to run"},
}
DEFAULT = "bge"

def spec(key):
    if key not in MODELS: raise KeyError(f"unknown embedding model '{key}'; choose from {', '.join(MODELS)}")
    return MODELS[key]