cached model list to memory (exclude falcon)

src/app.py

@@ -47,14 +47,24 @@ st.markdown(
         """,
         unsafe_allow_html=True,
     )
+
+@st.cache_resource()
+def cache_model_list():
+    model_list_info = {}
+    for model_name in model_list:
+        if not "tiiuae/falcon" in model_name: # Exclude Falcon models
+            model = get_model(model_name, library="transformers", access_token="")
+            model_list_info[model_name] = calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])
+    return model_list_info
+
 @st.cache_resource
 def get_gpu_specs():
     return pd.read_csv("data/gpu_specs.csv")
 
-@st.cache_resource
-def get_mistralai_table():
-    model = get_model("mistralai/Mistral-7B-v0.1", library="transformers", access_token="")
-    return calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])
+# @st.cache_resource
+# def get_mistralai_table():
+#     model = get_model("mistralai/Mistral-7B-v0.1", library="transformers", access_token="")
+#     return calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])
 
 def show_gpu_info(info, trainable_params=0, vendor=""):
     for var in ['Inference', 'Full Training Adam', 'LoRa Fine-tuning']:
@@ -91,7 +101,9 @@ def get_name(index):
 
 def custom_ceil(a, precision=0):
     return np.round(a + 0.5 * 10**(-precision), precision)
+
 gpu_specs = get_gpu_specs()
+model_list_info = cache_model_list()
 
 _, col, _ = st.columns([1,3,1])
 with col.expander("Information", expanded=True):
@@ -123,8 +135,8 @@ if model_name not in st.session_state:
         del st.session_state[st.session_state['actual_model']]
         del st.session_state['actual_model']
         gc.collect()
-    if model_name == "mistralai/Mistral-7B-v0.1": # cache Mistral
-        st.session_state[model_name] = get_mistralai_table()
+    if model_name in model_list_info.keys():
+        st.session_state[model_name] = model_list_info[model_name]
     else:
         model = get_model(model_name, library="transformers", access_token=access_token)
         st.session_state[model_name] = calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])