MODELS_ID = dict(distilroberta = 'flax-sentence-embeddings/st-codesearch-distilroberta-base',
                 mpnet = 'flax-sentence-embeddings/all_datasets_v3_mpnet-base',
                 minilm_l6 = 'flax-sentence-embeddings/all_datasets_v3_MiniLM-L6')

QA_MODELS_ID = dict(
    mpnet_asymmetric_qa = ['flax-sentence-embeddings/multi-QA_v1-mpnet-asymmetric-Q',
                           'flax-sentence-embeddings/multi-QA_v1-mpnet-asymmetric-A'],
    mpnet_qa='flax-sentence-embeddings/mpnet_stackexchange_v1',
    distilbert_qa = 'flax-sentence-embeddings/multi-qa_v1-distilbert-cls_dot'
)

SEARCH_MODELS_ID = dict(
    mpnet_qa='flax-sentence-embeddings/mpnet_stackexchange_v1'
)