import streamlit as st
from sentence_transformers import SentenceTransformer
from .config import MODELS_ID


@st.cache(allow_output_mutation=True)
def load_model(model_name):
    assert model_name in MODELS_ID.keys()
    # Lazy downloading
    model = SentenceTransformer(MODELS_ID[model_name])
    return model