Add file cache

app.py

@@ -1,3 +1,5 @@
+import os
+import pickle
 import streamlit as st
 import pandas as pd
 import vec2text
@@ -11,16 +13,36 @@ from sklearn.decomposition import PCA
 from streamlit_plotly_events import plotly_events
 import plotly.graph_objects as go
 import logging
+
 # Activate tqdm with pandas
 tqdm.pandas()
 
+# Custom file cache decorator
+def file_cache(file_path):
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            # Check if the file already exists
+            if os.path.exists(file_path):
+                # Load from cache
+                with open(file_path, "rb") as f:
+                    print(f"Loading cached data from {file_path}")
+                    return pickle.load(f)
+            else:
+                # Compute and save to cache
+                result = func(*args, **kwargs)
+                with open(file_path, "wb") as f:
+                    pickle.dump(result, f)
+                    print(f"Saving new cache to {file_path}")
+                return result
+        return wrapper
+    return decorator
+
 @st.cache_resource
 def vector_compressor_from_config():
-    'TODO'
-    # return PCA(n:n_components=2)
+    # Return UMAP with 2 components for dimensionality reduction
     return UMAP(n_components=2)
 
-# Caching the dataframe since loading from external source can be time-consuming
+# Caching the dataframe since loading from an external source can be time-consuming
 @st.cache_data
 def load_data():
     return pd.read_csv("https://huggingface.co/datasets/marksverdhei/reddit-syac-urls/resolve/main/train.csv")
@@ -50,8 +72,9 @@ def load_embeddings():
 
 embeddings = load_embeddings()
 
-# Caching UMAP reduction as it's a heavy computation
-@st.cache_resource
+# Custom cache the UMAP reduction using file_cache decorator
+@file_cache(".cache/reducer_embeddings.pickle")
+@st.cache_data
 def reduce_embeddings(embeddings):
     reducer = vector_compressor_from_config()
     return reducer.fit_transform(embeddings), reducer
@@ -79,15 +102,12 @@ fig.update_layout(
 # Display the scatterplot and capture click events
 selected_points = plotly_events(fig, click_event=True, hover_event=False, override_height=600, override_width="100%")
 
-
 # If a point is clicked, handle the embedding inversion
 if selected_points:
-    
     clicked_point = selected_points[0]
     x_coord = x = clicked_point['x']
     y_coord = y = clicked_point['y']
 
-
     inferred_embedding = reducer.inverse_transform(np.array([[x, y]]) if not isinstance(reducer, UMAP) else np.array([[x, y]]))
 
     inferred_embedding = inferred_embedding.astype("float32")
@@ -101,4 +121,4 @@ if selected_points:
     st.text(str(output))
     st.text(str(inferred_embedding))
 else:
-    st.text("Click on a point in the scatterplot to see its coordinates.")
+    st.text("Click on a point in the scatterplot to see its coordinates.")