pages/model_loader.py

import time
import streamlit as st
from streamlit.logger import get_logger
from utils.app_utils import is_model_alive
from app_config import ENDPOINT_NAMES

logger = get_logger(__name__)

st.set_page_config(page_title="Conversation Simulator")

models_alive = False
start = time.time()

MODELS2LOAD = {
    "CPC": {"model_name": "Phase Classifier", "loaded":None,},
    "CTL_llama3": {"model_name": "Texter Simulator", "loaded":None,},
    "BadPractices": {"model_name": "Advice Identificator", "loaded":None},
    "training_adherence": {"model_name": "Training Adherence", "loaded":None},
}

def write_model_status(writer, model_name, loaded, fail=False):
    if loaded == "200":
        writer.write(f"✅ - {model_name} Loaded")
    if fail:
        if loaded in ["400", "500"]:
            writer.write(f"❌ - {model_name} Failed to Load, Contact [email protected]")
        elif loaded == "404":
            writer.write(f"❌ - {model_name} Still loading, please try in a couple of minutes")
    else:
        writer.write(f"🔄 - {model_name} Loading")

with st.status("Loading Models Please Wait...(this may take up to 5 min)", expanded=True) as status:

    for k in MODELS2LOAD.keys():
        MODELS2LOAD[k]["writer"] = st.empty()

    while not models_alive:
        time.sleep(2)
        for name, config in MODELS2LOAD.items():
            config["loaded"] = is_model_alive(**ENDPOINT_NAMES[name])

        models_alive = all([x['loaded']=="200" for x in MODELS2LOAD.values()])
        
        for _, config in MODELS2LOAD.items():
            write_model_status(**config)
        
        if int(time.time()-start) > 30:
            status.update(
                label="Models took too long to load. Please Refresh Page in a couple of minutes", state="error", expanded=True
            )
            for _, config in MODELS2LOAD.items():
                write_model_status(**config, fail=True)
            break

if models_alive:
    st.switch_page("pages/convosim.py")