cache default model

src/app.py

@@ -22,6 +22,10 @@ st.markdown(
 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"])
 
 def show_gpu_info(info, trainable_params=0):
     for var in ['Inference', 'Full Training Adam', 'LoRa Fine-tuning']:
@@ -46,13 +50,6 @@ def get_name(index):
     row = gpu_specs.iloc[index]
     return f"{row['Product Name']} ({row['RAM (GB)']} GB, {row['Year']})"
 
-def create_plot(memory_table, y, title, container):
-    fig = px.bar(memory_table, x=memory_table.index, y=y, color_continuous_scale="RdBu_r")
-    fig.update_layout(yaxis_title="Number of GPUs", title=dict(text=title, font=dict(size=25)))
-    fig.update_coloraxes(showscale=False)
-    
-    container.plotly_chart(fig, use_container_width=True)
-
 gpu_specs = get_gpu_specs()
 
 access_token = st.sidebar.text_input("Access token")
@@ -61,16 +58,19 @@ if not model_name:
     st.info("Please enter a model name")
     st.stop()
 
-
-
 model_name = extract_from_url(model_name)
 if model_name not in st.session_state:
     if 'actual_model' in st.session_state:
         del st.session_state[st.session_state['actual_model']]
         del st.session_state['actual_model']
         gc.collect()
-    model = get_model(model_name, library="transformers", access_token=access_token)
-    st.session_state[model_name] = calculate_memory(model, ["float32", "float16/bfloat16", "int8", "int4"])
+    if model_name == "mistralai/Mistral-7B-v0.1": # cache Mistral
+        st.session_state[model_name] = get_mistralai_table()
+    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"])
+        del model
+        gc.collect()
     st.session_state['actual_model'] = model_name