| """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 = { |
| |
| "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] |
|
|
