backend/utils.py

import streamlit as st
from sentence_transformers import SentenceTransformer


@st.cache(allow_output_mutation=True)
def load_model(model_name, model_dict):
    assert model_name in model_dict.keys()
    # Lazy downloading
    model_ids = model_dict[model_name]
    if type(model_ids) == str:
        output = SentenceTransformer(model_ids)
    elif hasattr(model_ids, '__iter__'):
        output = [SentenceTransformer(name) for name in model_ids]

    return output