ok

app.py

@@ -1,9 +1,17 @@
 import streamlit as st
 import json
 
-import adrd
-ckpt_path = './ckpt_densenet_emb_encoder_2_AUPR.pt'
-mdl = adrd.model.ADRDModel.from_ckpt(ckpt_path, device='cpu')
+@st.cache(allow_output_mutation=True)
+def load_model():
+    import adrd
+    ckpt_path = './ckpt_densenet_emb_encoder_2_AUPR.pt'
+    model = adrd.model.ADRDModel.from_ckpt(ckpt_path, device='cpu')
+    return model
+
+def predict_proba(data_dict):
+    model = load_model()
+    pred_dict = model.predict_proba([data_dict])[1][0]
+    return pred_dict
 
 # Create a form for user input
 with st.form("json_input_form"):
@@ -16,7 +24,7 @@ if submit_button:
     try:
         # Parse the JSON input into a Python dictionary
         data_dict = json.loads(json_input)
-        pred_dict = mdl.predict_proba([data_dict])[1][0]
+        pred_dict = predict_proba(data_dict)
         st.write("Predicted probabilities:")
         st.json(pred_dict)
     except json.JSONDecodeError as e: