Change caching structure

pages/5_🗂_Organise_Demo.py

@@ -91,53 +91,69 @@ def generate_plot(X, y, filenames):
 
     return fig
 
-uploaded_files = st.sidebar.file_uploader("Select contracts to organise ", accept_multiple_files=True)
+@st.cache(allow_output_mutation=True)
+def prepare_figure(model, df):
+    X = [text[:500] for text in df['text'].to_list()]
+    filenames = df['filename'].to_list()
 
-button = st.sidebar.button('Organise Contracts', type='primary', use_container_width=True)
+    X_transform, y = get_transform_and_predictions(model, X)
 
-with st.container():
-    with st.spinner('⚙️ Loading model...'):
-        cuad_tfidf_umap_kmeans = load_model()
-        cuad_df = load_dataset()
+    fig = generate_plot(X_transform, y, filenames)
 
-        X = [text[:500] for text in cuad_df['text'].to_list()]
-        filenames = cuad_df['filename'].to_list()
+    return fig
 
-        X_transform, y = get_transform_and_predictions(cuad_tfidf_umap_kmeans, X)
+@st.cache()
+def prepare_page():
+    model = load_model()
+    df = load_dataset()
 
-        fig = generate_plot(X_transform, y, filenames)
+    X = [text[:500] for text in df['text'].to_list()]
+    filenames = df['filename'].to_list()
 
-        figure = st.plotly_chart(fig, use_container_width=True)
+    X_transform, y = get_transform_and_predictions(model, X)
 
-    if button:
-        figure.empty()
+    fig = prepare_figure(model, df)
 
-        with st.spinner('⚙️ Training model...'):
+    return fig, model
+
+
+uploaded_files = st.sidebar.file_uploader("Select contracts to organise ", accept_multiple_files=True)
 
-            if not uploaded_files or not len(uploaded_files) > 1:
-                st.write(
-                    "Please add at least two contracts"
-                )
+button = st.sidebar.button('Organise Contracts', type='primary', use_container_width=True)
+
+with st.spinner('⚙️ Loading model...'):
+    fig, cuad_tfidf_umap_kmeans = prepare_page()
+    figure = st.plotly_chart(fig, use_container_width=True)
+
+if button:
+    figure.empty()
+
+    with st.spinner('⚙️ Training model...'):
+
+        if not uploaded_files or not len(uploaded_files) > 2:
+            st.write(
+                "**Please add at least three contracts**"
+            )
+        else:
+            if len(uploaded_files) < 10:
+                n_clusters = 3
             else:
-                if len(uploaded_files) < 10:
-                    n_clusters = 3
-                else:
-                    n_clusters = 8
-            
-                X_train = [uploaded_file.read()[:500] for uploaded_file in uploaded_files]
-                filenames = [uploaded_file.name for uploaded_file in uploaded_files]
+                n_clusters = 8
+        
+            X_train = [uploaded_file.read()[:500] for uploaded_file in uploaded_files]
+            filenames = [uploaded_file.name for uploaded_file in uploaded_files]
 
-                tfidf_umap_kmeans = deepcopy(cuad_tfidf_umap_kmeans)
-                tfidf_umap_kmeans.set_params(kmeans__n_clusters=4)
-                tfidf_umap_kmeans.fit(X_train)
+            tfidf_umap_kmeans = deepcopy(cuad_tfidf_umap_kmeans)
+            tfidf_umap_kmeans.set_params(kmeans__n_clusters=n_clusters)
+            tfidf_umap_kmeans.fit(X_train)
 
-                X_transform, y = get_transform_and_predictions(cuad_tfidf_umap_kmeans, X_train)
+            X_transform, y = get_transform_and_predictions(cuad_tfidf_umap_kmeans, X_train)
 
-                fig = generate_plot(X_transform, y, filenames)
+            fig = generate_plot(X_transform, y, filenames)
 
-                st.write("**Your organised contracts:**")
+            st.write("**Your organised contracts:**")
 
-                st.plotly_chart(fig, use_container_width=True)
+            st.plotly_chart(fig, use_container_width=True)
             
 add_email_signup_form()