test: fix typo

app.py

@@ -2,196 +2,153 @@ import gzip
 import os
 import pickle
 from glob import glob
-from time import sleep
-
+import threading
+import psutil
 from functools import lru_cache
 import concurrent.futures
-from typing import Dict, Tuple, List
+from typing import Dict, Tuple, List, Optional
+from time import sleep
 
 import gradio as gr
 import numpy as np
-import plotly.graph_objects as go
 import torch
 from PIL import Image, ImageDraw
+import plotly.graph_objects as go
 from plotly.subplots import make_subplots
 
+# Constants
 IMAGE_SIZE = 400
 DATASET_LIST = ["imagenet", "oxford_flowers", "ucf101", "caltech101", "dtd", "eurosat"]
 GRID_NUM = 14
-pkl_root = "./data/out"
-preloaded_data = {}
-
-
-# Global cache for data
-_CACHE = {
-    'data_dict': {},
-    'sae_data_dict': {},
-    'model_data': {},
-    'segmasks': {},
-    'top_images': {},
-    'precomputed_activations' = {}
-
-}
+PKL_ROOT = "./data/out"
+
+# Global cache with better type hints and error handling
+class Cache:
+    def __init__(self):
+        self.data: Dict[str, Dict] = {
+            'data_dict': {},
+            'sae_data_dict': {},
+            'model_data': {},
+            'segmasks': {},
+            'top_images': {},
+            'precomputed_activations': {}
+        }
+        
+    def get(self, category: str, key: str, default=None):
+        try:
+            return self.data[category].get(key, default)
+        except KeyError:
+            return default
+            
+    def set(self, category: str, key: str, value):
+        try:
+            self.data[category][key] = value
+        except KeyError:
+            self.data[category] = {key: value}
+            
+    def clear_category(self, category: str):
+        if category in self.data:
+            self.data[category].clear()
+
+_CACHE = Cache()
 
 def load_all_data(image_root: str, pkl_root: str) -> Tuple[Dict, Dict]:
     """Load all data with optimized parallel processing."""
+    def load_image_file(image_file: str) -> Optional[Dict]:
+        try:
+            image = Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE))
+            return {
+                "image": image,
+                "image_path": image_file,
+            }
+        except Exception as e:
+            print(f"Error loading image {image_file}: {e}")
+            return None
+
     # Load images in parallel
     with concurrent.futures.ThreadPoolExecutor() as executor:
-        image_files = glob(f"{image_root}/*")
         future_to_file = {
-            executor.submit(_load_image_file, image_file): image_file 
-            for image_file in image_files
+            executor.submit(load_image_file, image_file): image_file 
+            for image_file in glob(f"{image_root}/*")
         }
         
         for future in concurrent.futures.as_completed(future_to_file):
-            image_file = future_to_file[future]
-            image_name = os.path.basename(image_file).split(".")[0]
-            result = future.result()
-            if result is not None:
-                _CACHE['data_dict'][image_name] = result
+            try:
+                image_file = future_to_file[future]
+                image_name = os.path.basename(image_file).split(".")[0]
+                result = future.result()
+                if result:
+                    _CACHE.set('data_dict', image_name, result)
+            except Exception as e:
+                print(f"Error processing image future: {e}")
 
     # Load SAE data
-    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
-        _CACHE['sae_data_dict']["mean_acts"] = pickle.load(f)
-
-    # Load mean act values in parallel
-    datasets = ["imagenet", "imagenet-sketch", "caltech101"]
-    _CACHE['sae_data_dict']["mean_act_values"] = {}
-    
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        future_to_dataset = {
-            executor.submit(_load_mean_act_values, dataset): dataset 
-            for dataset in datasets
-        }
-        
-        for future in concurrent.futures.as_completed(future_to_dataset):
-            dataset = future_to_dataset[future]
-            result = future.result()
-            if result is not None:
-                _CACHE['sae_data_dict']["mean_act_values"][dataset] = result
-
-    return _CACHE['data_dict'], _CACHE['sae_data_dict']
-
-def _load_image_file(image_file: str) -> Dict:
-    """Helper function to load a single image file."""
     try:
-        image = Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE))
-        return {
-            "image": image,
-            "image_path": image_file,
-        }
+        with open("./data/sae_data/mean_acts.pkl", "rb") as f:
+            _CACHE.set('sae_data_dict', "mean_acts", pickle.load(f))
     except Exception as e:
-        print(f"Error loading {image_file}: {e}")
-        return None
+        print(f"Error loading mean_acts.pkl: {e}")
 
-def _load_mean_act_values(dataset: str) -> np.ndarray:
-    """Helper function to load mean act values for a dataset."""
-    try:
-        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
-            return pickle.load(f)
-    except Exception as e:
-        print(f"Error loading mean act values for {dataset}: {e}")
-        return None
+    # Load mean act values
+    datasets = ["imagenet", "imagenet-sketch", "caltech101"]
+    for dataset in datasets:
+        try:
+            with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
+                if "mean_act_values" not in _CACHE.data['sae_data_dict']:
+                    _CACHE.set('sae_data_dict', "mean_act_values", {})
+                _CACHE.data['sae_data_dict']["mean_act_values"][dataset] = pickle.load(f)
+        except Exception as e:
+            print(f"Error loading mean act values for {dataset}: {e}")
+
+    return _CACHE.data['data_dict'], _CACHE.data['sae_data_dict']
 
 @lru_cache(maxsize=1024)
 def get_data(image_name: str, model_name: str) -> np.ndarray:
-    """Cached function to get model data."""
+    """Get model data with caching."""
     cache_key = f"{model_name}_{image_name}"
-    if cache_key not in _CACHE['model_data']:
-        data_dir = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
-        with gzip.open(data_dir, "rb") as f:
-            _CACHE['model_data'][cache_key] = pickle.load(f)
-    return _CACHE['model_data'][cache_key]
+    if cache_key not in _CACHE.data['model_data']:
+        try:
+            data_dir = f"{PKL_ROOT}/{model_name}/{image_name}.pkl.gz"
+            with gzip.open(data_dir, "rb") as f:
+                _CACHE.data['model_data'][cache_key] = pickle.load(f)
+        except Exception as e:
+            print(f"Error loading model data for {cache_key}: {e}")
+            return np.array([])
+    return _CACHE.data['model_data'][cache_key]
 
 @lru_cache(maxsize=1024)
 def get_activation_distribution(image_name: str, model_type: str) -> np.ndarray:
-    """Cached function to get activation distribution."""
-    activation = get_data(image_name, model_type)[0]
-    noisy_features_indices = (
-        (_CACHE['sae_data_dict']["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
-    )
-    activation[:, noisy_features_indices] = 0
-    return activation
-
-@lru_cache(maxsize=1024)
-def get_segmask(selected_image: str, slider_value: int, model_type: str) -> np.ndarray:
-    """Cached function to get segmentation mask."""
-    cache_key = f"{selected_image}_{slider_value}_{model_type}"
-    if cache_key not in _CACHE['segmasks']:
-        image = _CACHE['data_dict'][selected_image]["image"]
-        sae_act = get_data(selected_image, model_type)[0]
-        temp = sae_act[:, slider_value]
-        
-        mask = torch.Tensor(temp[1:].reshape(14, 14)).view(1, 1, 14, 14)
-        mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][0].numpy()
-        mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
-
-        base_opacity = 30
-        image_array = np.array(image)[..., :3]
-        rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
-        rgba_overlay[..., :3] = image_array[..., :3]
-
-        darkened_image = (image_array[..., :3] * (base_opacity / 255)).astype(np.uint8)
-        rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
-        rgba_overlay[..., 3] = 255
-
-        _CACHE['segmasks'][cache_key] = rgba_overlay
-        
-    return _CACHE['segmasks'][cache_key]
-
-@lru_cache(maxsize=1024)
-def get_top_images(slider_value: int, toggle_btn: bool) -> List[Image.Image]:
-    """Cached function to get top images."""
-    cache_key = f"{slider_value}_{toggle_btn}"
-    if cache_key not in _CACHE['top_images']:
-        dataset_path = "./data/top_images_masked" if toggle_btn else "./data/top_images"
-        paths = [
-            os.path.join(dataset_path, dataset, f"{slider_value}.jpg")
-            for dataset in ["imagenet", "imagenet-sketch", "caltech101"]
-        ]
-        
-        _CACHE['top_images'][cache_key] = [
-            Image.open(path) if os.path.exists(path) else Image.new("RGB", (256, 256), (255, 255, 255))
-            for path in paths
-        ]
+    """Get activation distribution with memory optimization."""
+    try:
+        activation = get_data(image_name, model_type)[0]
+        mean_acts = _CACHE.get('sae_data_dict', "mean_acts", {}).get("imagenet", np.array([]))
         
-    return _CACHE['top_images'][cache_key]
-
-
-# def preload_activation(image_name):
-#     for model in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
-#         image_file = f"{pkl_root}/{model}/{image_name}.pkl.gz"
-#         with gzip.open(image_file, "rb") as f:
-#             preloaded_data[model] = pickle.load(f)
-
-
-# def get_activation_distribution(image_name: str, model_type: str):
-#     activation = get_data(image_name, model_type)[0]
-
-#     noisy_features_indices = (
-#         (sae_data_dict["mean_acts"]["imagenet"] > 0.1).nonzero()[0].tolist()
-#     )
-#     activation[:, noisy_features_indices] = 0
-
-#     return activation
-
+        if mean_acts.size > 0:
+            noisy_features_indices = (mean_acts > 0.1).nonzero()[0]
+            activation[:, noisy_features_indices] = 0
+            
+        return activation
+    except Exception as e:
+        print(f"Error getting activation distribution: {e}")
+        return np.array([])
 
-def get_grid_loc(evt, image):
-    # Get click coordinates
+def get_grid_loc(evt: gr.EventData, image: Image.Image) -> Tuple[int, int, int, int]:
+    """Get grid location from click event."""
     x, y = evt._data["index"][0], evt._data["index"][1]
-
     cell_width = image.width // GRID_NUM
     cell_height = image.height // GRID_NUM
-
     grid_x = x // cell_width
     grid_y = y // cell_height
     return grid_x, grid_y, cell_width, cell_height
 
-
-def highlight_grid(evt: gr.EventData, image_name):
-    image = data_dict[image_name]["image"]
+def highlight_grid(evt: gr.EventData, image_name: str) -> Image.Image:
+    """Highlight selected grid cell."""
+    image = _CACHE.get('data_dict', image_name, {}).get("image")
+    if not image:
+        return None
+        
     grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
-
+    
     highlighted_image = image.copy()
     draw = ImageDraw.Draw(highlighted_image)
     box = [
@@ -201,25 +158,18 @@ def highlight_grid(evt: gr.EventData, image_name):
         (grid_y + 1) * cell_height,
     ]
     draw.rectangle(box, outline="red", width=3)
-
     return highlighted_image
 
-
-def load_image(img_name):
-    return Image.open(data_dict[img_name]["image_path"]).resize(
-        (IMAGE_SIZE, IMAGE_SIZE)
-    )
-
-
 def plot_activations(
-    all_activation,
-    tile_activations=None,
-    grid_x=None,
-    grid_y=None,
-    top_k=5,
-    colors=("blue", "cyan"),
-    model_name="CLIP",
-):
+    all_activation: np.ndarray,
+    tile_activations: Optional[np.ndarray] = None,
+    grid_x: Optional[int] = None,
+    grid_y: Optional[int] = None,
+    top_k: int = 5,
+    colors: Tuple[str, str] = ("blue", "cyan"),
+    model_name: str = "CLIP",
+) -> go.Figure:
+    """Plot activation distributions."""
     fig = go.Figure()
 
     def _add_scatter_with_annotation(fig, activations, model_name, color, label):
@@ -248,59 +198,90 @@ def plot_activations(
             )
         return fig
 
-    label = f"{model_name.split('-')[-0]} Image-level"
-    fig = _add_scatter_with_annotation(
-        fig, all_activation, model_name, colors[0], label
-    )
+    label = f"{model_name.split('-')[-1]} Image-level"
+    fig = _add_scatter_with_annotation(fig, all_activation, model_name, colors[0], label)
+    
     if tile_activations is not None:
-        label = f"{model_name.split('-')[-0]} Tile ({grid_x}, {grid_y})"
-        fig = _add_scatter_with_annotation(
-            fig, tile_activations, model_name, colors[1], label
-        )
+        label = f"{model_name.split('-')[-1]} Tile ({grid_x}, {grid_y})"
+        fig = _add_scatter_with_annotation(fig, tile_activations, model_name, colors[1], label)
 
     fig.update_layout(
         title="Activation Distribution",
         xaxis_title="SAE latent index",
         yaxis_title="Activation Value",
         template="plotly_white",
-    )
-    fig.update_layout(
         legend=dict(orientation="h", yanchor="middle", y=0.5, xanchor="center", x=0.5)
     )
 
     return fig
 
+def get_segmask(selected_image: str, slider_value: int, model_type: str) -> Optional[np.ndarray]:
+    """Get segmentation mask with caching."""
+    cache_key = f"{selected_image}_{slider_value}_{model_type}"
+    cached_mask = _CACHE.get('segmasks', cache_key)
+    if cached_mask is not None:
+        return cached_mask
 
-def get_activations(evt: gr.EventData, selected_image: str, model_name: str, colors):
-    activation = get_activation_distribution(selected_image, model_name)
-    all_activation = activation.mean(0)
-
-    tile_activations = None
-    grid_x = None
-    grid_y = None
-
-    if evt is not None:
-        if evt._data is not None:
-            image = data_dict[selected_image]["image"]
-            grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
-            token_idx = grid_y * GRID_NUM + grid_x + 1
-            tile_activations = activation[token_idx]
-
-    fig = plot_activations(
-        all_activation,
-        tile_activations,
-        grid_x,
-        grid_y,
-        top_k=5,
-        model_name=model_name,
-        colors=colors,
-    )
-    return fig
+    try:
+        image = _CACHE.get('data_dict', selected_image, {}).get("image")
+        if image is None:
+            return None
+
+        sae_act = get_data(selected_image, model_type)[0]
+        temp = sae_act[:, slider_value]
+        
+        mask = torch.tensor(temp[1:].reshape(14, 14)).view(1, 1, 14, 14)
+        mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][0].numpy()
+        
+        if mask.size == 0:
+            return None
+            
+        mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
+        
+        base_opacity = 30
+        image_array = np.array(image)[..., :3]
+        rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+        rgba_overlay[..., :3] = image_array
+        
+        darkened_image = (image_array * (base_opacity / 255)).astype(np.uint8)
+        rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
+        rgba_overlay[..., 3] = 255
 
+        _CACHE.set('segmasks', cache_key, rgba_overlay)
+        return rgba_overlay
+        
+    except Exception as e:
+        print(f"Error generating segmentation mask: {e}")
+        return None
+
+def get_top_images(slider_value: int, toggle_btn: bool) -> List[Image.Image]:
+    """Get top images with caching."""
+    cache_key = f"{slider_value}_{toggle_btn}"
+    cached_images = _CACHE.get('top_images', cache_key)
+    if cached_images is not None:
+        return cached_images
+
+    dataset_path = "./data/top_images_masked" if toggle_btn else "./data/top_images"
+    paths = [
+        os.path.join(dataset_path, dataset, f"{slider_value}.jpg")
+        for dataset in ["imagenet", "imagenet-sketch", "caltech101"]
+    ]
+    
+    images = [
+        Image.open(path) if os.path.exists(path) else Image.new("RGB", (256, 256), (255, 255, 255))
+        for path in paths
+    ]
+    
+    _CACHE.set('top_images', cache_key, images)
+    return images
 
+# UI Event Handlers
 def plot_activation_distribution(
-    evt: gr.EventData, selected_image: str, model_name: str
-):
+    evt: Optional[gr.EventData],
+    selected_image: str,
+    model_name: str
+) -> go.Figure:
+    """Plot activation distributions for both models."""
     fig = make_subplots(
         rows=2,
         cols=1,
@@ -308,17 +289,37 @@ def plot_activation_distribution(
         subplot_titles=["CLIP Activation", f"{model_name} Activation"],
     )
 
-    fig_clip = get_activations(
-        evt, selected_image, "CLIP", colors=("#00b4d8", "#90e0ef")
-    )
-    fig_maple = get_activations(
-        evt, selected_image, model_name, colors=("#ff5a5f", "#ffcad4")
-    )
+    def get_activations(evt, selected_image, model_name, colors):
+        activation = get_activation_distribution(selected_image, model_name)
+        all_activation = activation.mean(0)
+
+        tile_activations = None
+        grid_x = None
+        grid_y = None
+
+        if evt is not None and evt._data is not None:
+            image = _CACHE.get('data_dict', selected_image, {}).get("image")
+            if image:
+                grid_x, grid_y, _, _ = get_grid_loc(evt, image)
+                token_idx = grid_y * GRID_NUM + grid_x + 1
+                tile_activations = activation[token_idx]
+
+        return plot_activations(
+            all_activation,
+            tile_activations,
+            grid_x,
+            grid_y,
+            top_k=5,
+            model_name=model_name,
+            colors=colors,
+        )
+
+    fig_clip = get_activations(evt, selected_image, "CLIP", colors=("#00b4d8", "#90e0ef"))
+    fig_maple = get_activations(evt, selected_image, model_name, colors=("#ff5a5f", "#ffcad4"))
 
     def _attach_fig(fig, sub_fig, row, col, yref):
         for trace in sub_fig.data:
             fig.add_trace(trace, row=row, col=col)
-
         for annotation in sub_fig.layout.annotations:
             annotation.update(yref=yref)
             fig.add_annotation(annotation)
@@ -332,8 +333,6 @@ def plot_activation_distribution(
     fig.update_yaxes(title_text="Activation Value", row=1, col=1)
     fig.update_yaxes(title_text="Activation Value", row=2, col=1)
     fig.update_layout(
-        # height=500,
-        # title="Activation Distributions",
         template="plotly_white",
         showlegend=True,
         legend=dict(orientation="h", yanchor="bottom", y=-0.2, xanchor="center", x=0.5),
@@ -342,73 +341,12 @@ def plot_activation_distribution(
 
     return fig
 
-
-# def get_segmask(selected_image, slider_value, model_type):
-#     image = data_dict[selected_image]["image"]
-#     sae_act = get_data(selected_image, model_type)[0]
-#     temp = sae_act[:, slider_value]
-#     try:
-#         mask = torch.Tensor(temp[1:,].reshape(14, 14)).view(1, 1, 14, 14)
-#     except Exception as e:
-#         print(sae_act.shape, slider_value)
-#     mask = torch.nn.functional.interpolate(mask, (image.height, image.width))[0][
-#         0
-#     ].numpy()
-#     mask = (mask - mask.min()) / (mask.max() - mask.min() + 1e-10)
-
-#     base_opacity = 30
-#     image_array = np.array(image)[..., :3]
-#     rgba_overlay = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
-#     rgba_overlay[..., :3] = image_array[..., :3]
-
-#     darkened_image = (image_array[..., :3] * (base_opacity / 255)).astype(np.uint8)
-#     rgba_overlay[mask == 0, :3] = darkened_image[mask == 0]
-#     rgba_overlay[..., 3] = 255  # Fully opaque
-
-#     return rgba_overlay
-
-
-# def get_top_images(slider_value, toggle_btn):
-#     def _get_images(dataset_path):
-#         top_image_paths = [
-#             os.path.join(dataset_path, "imagenet", f"{slider_value}.jpg"),
-#             os.path.join(dataset_path, "imagenet-sketch", f"{slider_value}.jpg"),
-#             os.path.join(dataset_path, "caltech101", f"{slider_value}.jpg"),
-#         ]
-#         top_images = [
-#             (
-#                 Image.open(path)
-#                 if os.path.exists(path)
-#                 else Image.new("RGB", (256, 256), (255, 255, 255))
-#             )
-#             for path in top_image_paths
-#         ]
-#         return top_images
-
-#     if toggle_btn:
-#         top_images = _get_images("./data/top_images_masked")
-#     else:
-#         top_images = _get_images("./data/top_images")
-#     return top_images
-
-
-def show_activation_heatmap(selected_image, slider_value, model_type, toggle_btn=False):
-    slider_value = int(slider_value.split("-")[-1])
-    rgba_overlay = get_segmask(selected_image, slider_value, model_type)
-    top_images = get_top_images(slider_value, toggle_btn)
-
-    act_values = []
-    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
-        act_value = sae_data_dict["mean_act_values"][dataset][slider_value, :5]
-        act_value = [str(round(value, 3)) for value in act_value]
-        act_value = " | ".join(act_value)
-        out = f"#### Activation values: {act_value}"
-        act_values.append(out)
-
-    return rgba_overlay, top_images, act_values
-
-
-def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
+def show_activation_heatmap_clip(
+    selected_image: str,
+    slider_value: str,
+    toggle_btn: bool
+):
+    """Show activation heatmap for CLIP model."""
     rgba_overlay, top_images, act_values = show_activation_heatmap(
         selected_image, slider_value, "CLIP", toggle_btn
     )
@@ -423,49 +361,68 @@ def show_activation_heatmap_clip(selected_image, slider_value, toggle_btn):
         act_values[2],
     )
 
+def show_activation_heatmap(
+    selected_image: str,
+    slider_value: str,
+    model_type: str,
+    toggle_btn: bool = False
+) -> Tuple[np.ndarray, List[Image.Image], List[str]]:
+    """Show activation heatmap with segmentation mask and top images."""
+    slider_value = int(slider_value.split("-")[-1])
+    rgba_overlay = get_segmask(selected_image, slider_value, model_type)
+    top_images = get_top_images(slider_value, toggle_btn)
+
+    act_values = []
+    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
+        act_value = _CACHE.get('sae_data_dict', "mean_act_values", {}).get(dataset, np.array([]))[slider_value, :5]
+        act_value = [str(round(value, 3)) for value in act_value]
+        act_value = " | ".join(act_value)
+        out = f"#### Activation values: {act_value}"
+        act_values.append(out)
 
-def show_activation_heatmap_maple(selected_image, slider_value, model_name):
+    return rgba_overlay, top_images, act_values
+
+def show_activation_heatmap_maple(
+    selected_image: str,
+    slider_value: str,
+    model_name: str
+) -> np.ndarray:
+    """Show activation heatmap for MaPLE model."""
     slider_value = int(slider_value.split("-")[-1])
     rgba_overlay = get_segmask(selected_image, slider_value, model_name)
     sleep(0.1)
     return rgba_overlay
 
-
-def get_init_radio_options(selected_image, model_name):
+def get_init_radio_options(selected_image: str, model_name: str) -> List[str]:
+    """Get initial radio options for UI."""
     clip_neuron_dict = {}
     maple_neuron_dict = {}
 
-    def _get_top_actvation(selected_image, model_name, neuron_dict, top_k=5):
+    def _get_top_activation(selected_image: str, model_name: str, neuron_dict: Dict, top_k: int = 5) -> Dict:
         activations = get_activation_distribution(selected_image, model_name).mean(0)
         top_neurons = list(np.argsort(activations)[::-1][:top_k])
         for top_neuron in top_neurons:
             neuron_dict[top_neuron] = activations[top_neuron]
-        sorted_dict = dict(
-            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
-        )
-        return sorted_dict
-
-    clip_neuron_dict = _get_top_actvation(selected_image, "CLIP", clip_neuron_dict)
-    maple_neuron_dict = _get_top_actvation(
-        selected_image, model_name, maple_neuron_dict
-    )
-
-    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
+        return dict(sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True))
 
-    return radio_choices
+    clip_neuron_dict = _get_top_activation(selected_image, "CLIP", clip_neuron_dict)
+    maple_neuron_dict = _get_top_activation(selected_image, model_name, maple_neuron_dict)
 
+    return get_radio_names(clip_neuron_dict, maple_neuron_dict)
 
-def get_radio_names(clip_neuron_dict, maple_neuron_dict):
+def get_radio_names(
+    clip_neuron_dict: Dict[int, float],
+    maple_neuron_dict: Dict[int, float]
+) -> List[str]:
+    """Generate radio button names based on neuron activations."""
     clip_keys = list(clip_neuron_dict.keys())
     maple_keys = list(maple_neuron_dict.keys())
 
     common_keys = list(set(clip_keys).intersection(set(maple_keys)))
-    clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
-    maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
+    clip_only_keys = list(set(clip_keys) - set(maple_keys))
+    maple_only_keys = list(set(maple_keys) - set(clip_keys))
 
-    common_keys.sort(
-        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
-    )
+    common_keys.sort(key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True)
     clip_only_keys.sort(reverse=True)
     maple_only_keys.sort(reverse=True)
 
@@ -476,81 +433,54 @@ def get_radio_names(clip_neuron_dict, maple_neuron_dict):
 
     return out
 
+def update_radio_options(
+    evt: Optional[gr.EventData],
+    selected_image: str,
+    model_name: str
+) -> gr.Radio:
+    """Update radio options based on user interaction."""
+    clip_neuron_dict = {}
+    maple_neuron_dict = {}
 
-def update_radio_options(evt: gr.EventData, selected_image, model_name):
-    def _sort_and_save_top_k(activations, neuron_dict, top_k=5):
-        top_neurons = list(np.argsort(activations)[::-1][:top_k])
-        for top_neuron in top_neurons:
-            neuron_dict[top_neuron] = activations[top_neuron]
-
-    def _get_top_actvation(evt, selected_image, model_name, neuron_dict):
+    def _get_top_activation(evt, selected_image, model_name, neuron_dict):
         all_activation = get_activation_distribution(selected_image, model_name)
         image_activation = all_activation.mean(0)
-        _sort_and_save_top_k(image_activation, neuron_dict)
+        top_neurons = list(np.argsort(image_activation)[::-1][:5])
+        for top_neuron in top_neurons:
+            neuron_dict[top_neuron] = image_activation[top_neuron]
 
-        if evt is not None:
-            if evt._data is not None and isinstance(evt._data["index"], list):
-                image = data_dict[selected_image]["image"]
-                grid_x, grid_y, cell_width, cell_height = get_grid_loc(evt, image)
+        if evt is not None and evt._data is not None and isinstance(evt._data["index"], list):
+            image = _CACHE.get('data_dict', selected_image, {}).get("image")
+            if image:
+                grid_x, grid_y, _, _ = get_grid_loc(evt, image)
                 token_idx = grid_y * GRID_NUM + grid_x + 1
                 tile_activations = all_activation[token_idx]
-                _sort_and_save_top_k(tile_activations, neuron_dict)
-
-        sorted_dict = dict(
-            sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True)
-        )
-        return sorted_dict
-
-    clip_neuron_dict = {}
-    maple_neuron_dict = {}
-    clip_neuron_dict = _get_top_actvation(evt, selected_image, "CLIP", clip_neuron_dict)
-    maple_neuron_dict = _get_top_actvation(
-        evt, selected_image, model_name, maple_neuron_dict
-    )
+                top_tile_neurons = list(np.argsort(tile_activations)[::-1][:5])
+                for top_neuron in top_tile_neurons:
+                    neuron_dict[top_neuron] = tile_activations[top_neuron]
 
-    clip_keys = list(clip_neuron_dict.keys())
-    maple_keys = list(maple_neuron_dict.keys())
-
-    common_keys = list(set(clip_keys).intersection(set(maple_keys)))
-    clip_only_keys = list(set(clip_keys) - (set(maple_keys)))
-    maple_only_keys = list(set(maple_keys) - (set(clip_keys)))
-
-    common_keys.sort(
-        key=lambda x: max(clip_neuron_dict[x], maple_neuron_dict[x]), reverse=True
-    )
-    clip_only_keys.sort(reverse=True)
-    maple_only_keys.sort(reverse=True)
+        return dict(sorted(neuron_dict.items(), key=lambda item: item[1], reverse=True))
 
-    out = []
-    out.extend([f"common-{i}" for i in common_keys[:5]])
-    out.extend([f"CLIP-{i}" for i in clip_only_keys[:5]])
-    out.extend([f"MaPLE-{i}" for i in maple_only_keys[:5]])
-
-    radio_choices = gr.Radio(
-        choices=out, label="Top activating SAE latent", value=out[0]
-    )
-    sleep(0.1)
-    return radio_choices
+    clip_neuron_dict = _get_top_activation(evt, selected_image, "CLIP", clip_neuron_dict)
+    maple_neuron_dict = _get_top_activation(evt, selected_image, model_name, maple_neuron_dict)
 
+    radio_choices = get_radio_names(clip_neuron_dict, maple_neuron_dict)
+    return gr.Radio(choices=radio_choices, label="Top activating SAE latent", value=radio_choices[0])
 
-def update_markdown(option_value):
+def update_markdown(option_value: str) -> Tuple[str, str]:
+    """Update markdown text based on selected option."""
     latent_idx = int(option_value.split("-")[-1])
     out_1 = f"## Segmentation mask for the selected SAE latent - {latent_idx}"
     out_2 = f"## Top reference images for the selected SAE latent - {latent_idx}"
     return out_1, out_2
 
-
-def get_data(image_name, model_name):
-    pkl_root = "./data/out"
-    data_dir = f"{pkl_root}/{model_name}/{image_name}.pkl.gz"
-    with gzip.open(data_dir, "rb") as f:
-        data = pickle.load(f)
-        out = data
-
-    return out
-
-
-def update_all(selected_image, slider_value, toggle_btn, model_name):
+def update_all(
+    selected_image: str,
+    slider_value: str,
+    toggle_btn: bool,
+    model_name: str
+) -> Tuple:
+    """Update all UI components."""
     (
         seg_mask_display,
         top_image_1,
@@ -560,6 +490,7 @@ def update_all(selected_image, slider_value, toggle_btn, model_name):
         act_value_2,
         act_value_3,
     ) = show_activation_heatmap_clip(selected_image, slider_value, toggle_btn)
+    
     seg_mask_display_maple = show_activation_heatmap_maple(
         selected_image, slider_value, model_name
     )
@@ -578,101 +509,67 @@ def update_all(selected_image, slider_value, toggle_btn, model_name):
         markdown_display_2,
     )
 
-
-def load_all_data(image_root, pkl_root):
-    image_files = glob(f"{image_root}/*")
-    data_dict = {}
-    for image_file in image_files:
-        image_name = os.path.basename(image_file).split(".")[0]
-        if image_file not in data_dict:
-            data_dict[image_name] = {
-                "image": Image.open(image_file).resize((IMAGE_SIZE, IMAGE_SIZE)),
-                "image_path": image_file,
-            }
-
-    sae_data_dict = {}
-    with open("./data/sae_data/mean_acts.pkl", "rb") as f:
-        data = pickle.load(f)
-        sae_data_dict["mean_acts"] = data
-
-    sae_data_dict["mean_act_values"] = {}
-    for dataset in ["imagenet", "imagenet-sketch", "caltech101"]:
-        with gzip.open(f"./data/sae_data/mean_act_values_{dataset}.pkl.gz", "rb") as f:
-            data = pickle.load(f)
-            sae_data_dict["mean_act_values"][dataset] = data
-
-    return data_dict, sae_data_dict
-
-
-def preload_all_model_data():
-    """Preload all model data into memory at startup"""
-    print("Preloading model data...")
-    for image_name in data_dict.keys():
-        for model_name in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
-            try:
-                data = get_data(image_name, model_name)
-                cache_key = f"{model_name}_{image_name}"
-                _CACHE['model_data'][cache_key] = data
-            except Exception as e:
-                print(f"Error preloading {cache_key}: {e}")
-
-def precompute_activations():
-    """Precompute and cache common activation patterns"""
-    print("Precomputing activations...")
-    for image_name in data_dict.keys():
-        for model_name in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
-            activation = get_activation_distribution(image_name, model_name)
-            cache_key = f"activation_{model_name}_{image_name}"
-            _CACHE['precomputed_activations'][cache_key] = activation.mean(0)
-
-
-
-def precompute_segmasks():
-    """Precompute common segmentation masks"""
-    print("Precomputing segmentation masks...")
-    for image_name in data_dict.keys():
-        for model_type in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
-            for slider_value in range(0, 100):  # Adjust range as needed
-                try:
-                    mask = get_segmask(image_name, slider_value, model_type)
-                    cache_key = f"{image_name}_{slider_value}_{model_type}"
-                    _CACHE['segmasks'][cache_key] = mask
-                except Exception as e:
-                    print(f"Error precomputing mask {cache_key}: {e}")
-
-data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=pkl_root)
+def monitor_memory_usage():
+    """Monitor memory usage and clean cache if necessary."""
+    process = psutil.Process()
+    mem_info = process.memory_info()
+    mem_percent = process.memory_percent()
+    
+    print(f"""
+    Memory Usage:
+    - RSS: {mem_info.rss / (1024**2):.2f} MB
+    - VMS: {mem_info.vms / (1024**2):.2f} MB
+    - Percent: {mem_percent:.1f}%
+    - Cache Sizes: {[len(cache) for cache in _CACHE.data.values()]}
+    """)
+    
+    if mem_percent > 80:
+        print("Memory usage too high, clearing caches...")
+        _CACHE.clear_category('segmasks')
+        _CACHE.clear_category('top_images')
+        _CACHE.clear_category('precomputed_activations')
+
+def start_memory_monitor(interval: int = 300):
+    """Start periodic memory monitoring."""
+    monitor_memory_usage()
+    threading.Timer(interval, start_memory_monitor).start()
+
+# Initialize the application
+data_dict, sae_data_dict = load_all_data(image_root="./data/image", pkl_root=PKL_ROOT)
 default_image_name = "christmas-imagenet"
 
+
+# Create the Gradio interface
 with gr.Blocks(
     theme=gr.themes.Citrus(),
     css="""
     .image-row .gr-image { margin: 0 !important; padding: 0 !important; }
-    .image-row img { width: auto; height: 50px; } /* Set a uniform height for all images */
+    .image-row img { width: auto; height: 50px; }
 """,
 ) as demo:
     with gr.Row():
         with gr.Column():
-            # Left View: Image selection and click handling
             gr.Markdown("## Select input image and patch on the image")
             image_selector = gr.Dropdown(
-                choices=list(data_dict.keys()),
+                choices=list(_CACHE.data['data_dict'].keys()),
                 value=default_image_name,
                 label="Select Image",
             )
             image_display = gr.Image(
-                value=data_dict[default_image_name]["image"],
+                value=_CACHE.get('data_dict', default_image_name, {}).get("image"),
                 type="pil",
                 interactive=True,
             )
 
-            # Update image display when a new image is selected
             image_selector.change(
-                fn=lambda img_name: data_dict[img_name]["image"],
+                fn=lambda img_name: _CACHE.get('data_dict', img_name, {}).get("image"),
                 inputs=image_selector,
                 outputs=image_display,
             )
             image_display.select(
-                fn=highlight_grid, inputs=[image_selector], outputs=[image_display]
+                fn=highlight_grid,
+                inputs=[image_selector],
+                outputs=[image_display]
             )
 
         with gr.Column():
@@ -683,12 +580,8 @@ with gr.Blocks(
                 value=model_options[0],
                 label="Select adapted model (MaPLe)",
             )
-            init_plot = plot_activation_distribution(
-                None, default_image_name, model_options[0]
-            )
-            neuron_plot = gr.Plot(
-                label="Neuron Activation", value=init_plot, show_label=False
-            )
+            init_plot = plot_activation_distribution(None, default_image_name, model_options[0])
+            neuron_plot = gr.Plot(value=init_plot, show_label=False)
 
             image_selector.change(
                 fn=plot_activation_distribution,
@@ -701,7 +594,9 @@ with gr.Blocks(
                 outputs=neuron_plot,
             )
             model_selector.change(
-                fn=load_image, inputs=[image_selector], outputs=image_display
+                fn=lambda img_name: _CACHE.get('data_dict', img_name, {}).get("image"),
+                inputs=[image_selector],
+                outputs=image_display,
             )
             model_selector.change(
                 fn=plot_activation_distribution,
@@ -712,10 +607,9 @@ with gr.Blocks(
     with gr.Row():
         with gr.Column():
             radio_names = get_init_radio_options(default_image_name, model_options[0])
-
-            feautre_idx = radio_names[0].split("-")[-1]
+            feature_idx = radio_names[0].split("-")[-1]
             markdown_display = gr.Markdown(
-                f"## Segmentation mask for the selected SAE latent - {feautre_idx}"
+                f"## Segmentation mask for the selected SAE latent - {feature_idx}"
             )
             init_seg, init_tops, init_values = show_activation_heatmap(
                 default_image_name, radio_names[0], "CLIP"
@@ -727,13 +621,10 @@ with gr.Blocks(
                 default_image_name, radio_names[0], model_options[0]
             )
             gr.Markdown("### Localize SAE latent activation using MaPLE")
-            seg_mask_display_maple = gr.Image(
-                value=init_seg_maple, type="pil", show_label=False
-            )
+            seg_mask_display_maple = gr.Image(value=init_seg_maple, type="pil", show_label=False)
 
         with gr.Column():
             gr.Markdown("## Top activating SAE latent index")
-
             radio_choices = gr.Radio(
                 choices=radio_names,
                 label="Top activating SAE latent",
@@ -741,144 +632,106 @@ with gr.Blocks(
                 value=radio_names[0],
             )
             toggle_btn = gr.Checkbox(label="Show segmentation mask", value=False)
-
             markdown_display_2 = gr.Markdown(
-                f"## Top reference images for the selected SAE latent - {feautre_idx}"
+                f"## Top reference images for the selected SAE latent - {feature_idx}"
             )
 
             gr.Markdown("### ImageNet")
-            top_image_1 = gr.Image(
-                value=init_tops[0], type="pil", label="ImageNet", show_label=False
-            )
+            top_image_1 = gr.Image(value=init_tops[0], type="pil", show_label=False)
             act_value_1 = gr.Markdown(init_values[0])
 
             gr.Markdown("### ImageNet-Sketch")
-            top_image_2 = gr.Image(
-                value=init_tops[1],
-                type="pil",
-                label="ImageNet-Sketch",
-                show_label=False,
-            )
+            top_image_2 = gr.Image(value=init_tops[1], type="pil", show_label=False)
             act_value_2 = gr.Markdown(init_values[1])
 
             gr.Markdown("### Caltech101")
-            top_image_3 = gr.Image(
-                value=init_tops[2], type="pil", label="Caltech101", show_label=False
-            )
+            top_image_3 = gr.Image(value=init_tops[2], type="pil", show_label=False)
             act_value_3 = gr.Markdown(init_values[2])
 
+            # Event handlers
             image_display.select(
                 fn=update_radio_options,
                 inputs=[image_selector, model_selector],
                 outputs=[radio_choices],
             )
-
             model_selector.change(
                 fn=update_radio_options,
                 inputs=[image_selector, model_selector],
                 outputs=[radio_choices],
             )
-
             image_selector.select(
                 fn=update_radio_options,
                 inputs=[image_selector, model_selector],
                 outputs=[radio_choices],
             )
+            radio_choices.change(
+                fn=update_all,
+                inputs=[image_selector, radio_choices, toggle_btn, model_selector],
+                outputs=[
+                    seg_mask_display,
+                    seg_mask_display_maple,
+                    top_image_1,
+                    top_image_2,
+                    top_image_3,
+                    act_value_1,
+                    act_value_2,
+                    act_value_3,
+                    markdown_display,
+                    markdown_display_2,
+                ],
+            )
 
-        radio_choices.change(
-            fn=update_all,
-            inputs=[image_selector, radio_choices, toggle_btn, model_selector],
-            outputs=[
-                seg_mask_display,
-                seg_mask_display_maple,
-                top_image_1,
-                top_image_2,
-                top_image_3,
-                act_value_1,
-                act_value_2,
-                act_value_3,
-                markdown_display,
-                markdown_display_2,
-            ],
-        )
-
-        toggle_btn.change(
-            fn=show_activation_heatmap_clip,
-            inputs=[image_selector, radio_choices, toggle_btn],
-            outputs=[
-                seg_mask_display,
-                top_image_1,
-                top_image_2,
-                top_image_3,
-                act_value_1,
-                act_value_2,
-                act_value_3,
-            ],
-        )
-
-    # Launch the app
-    # demo.queue()
-    # demo.launch()
-
-    
-# if __name__ == "__main__":
-#     demo.queue()  # Enable queuing for better handling of concurrent users
-#     demo.launch(
-#         server_name="0.0.0.0",  # Allow external access
-#         server_port=7860,
-#         share=False,  # Set to True if you want to create a public URL
-#         show_error=True,
-#         # Optimize concurrency
-#         max_threads=8,  # Adjust based on your CPU cores
-#     )
+            toggle_btn.change(
+                fn=show_activation_heatmap_clip,
+                inputs=[image_selector, radio_choices, toggle_btn],
+                outputs=[
+                    seg_mask_display,
+                    top_image_1,
+                    top_image_2,
+                    top_image_3,
+                    act_value_1,
+                    act_value_2,
+                    act_value_3,
+                ],
+            )
 
 if __name__ == "__main__":
-    import psutil
+    # Initialize memory monitoring
+    start_memory_monitor()
     
     # Get system memory info
     mem = psutil.virtual_memory()
     total_ram_gb = mem.total / (1024**3)
     
-    # Configure cache sizes based on available RAM
-    cache_size = int(total_ram_gb * 100)  # Rough estimate: 100 entries per GB
-
-    # Precompute all data
-    print("Starting precomputation...")
-    preload_all_model_data()
-    precompute_activations()
-    precompute_segmasks()
-    print("Precomputation complete!")
-    
-    # Memory monitoring function
-    def monitor_memory_usage():
-        """Monitor and log memory usage"""
-        process = psutil.Process()
-        mem_info = process.memory_info()
-        print(f"""
-        Memory Usage:
-        - RSS: {mem_info.rss / (1024**2):.2f} MB
-        - VMS: {mem_info.vms / (1024**2):.2f} MB
-        - Cache Size: {len(_CACHE['model_data'])} entries
-        """)
-
-    # Start periodic monitoring
-    def start_memory_monitor():
-        threading.Timer(300.0, start_memory_monitor).start()  # Every 5 minutes
-        monitor_memory_usage()
-
-    # Start the monitoring
-    import threading
-    start_memory_monitor()
-
-    # Launch the app with memory-optimized settings
-    demo.queue(max_size=min(20, int(total_ram_gb)))  # Scale queue with RAM
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_error=True,
-        max_threads=min(16, psutil.cpu_count()),  # Scale threads with CPU
-        websocket_ping_timeout=60,
-        preventive_refresh=True,
-        memory_limit_mb=int(total_ram_gb * 1024 * 0.8)  # Use up to 80% of RAM
-    )
+    try:
+        print("Starting application initialization...")
+        
+        # Precompute common data
+        print("Precomputing activation patterns...")
+        for image_name in _CACHE.data['data_dict'].keys():
+            for model_name in ["CLIP"] + [f"MaPLE-{ds}" for ds in DATASET_LIST]:
+                try:
+                    activation = get_activation_distribution(image_name, model_name)
+                    cache_key = f"activation_{model_name}_{image_name}"
+                    _CACHE.set('precomputed_activations', cache_key, activation.mean(0))
+                except Exception as e:
+                    print(f"Error precomputing activation for {image_name}, {model_name}: {e}")
+
+        print("Starting Gradio interface...")
+        # Launch the app with optimized settings
+        demo.queue(max_size=min(20, int(total_ram_gb)))
+        demo.launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            share=False,
+            show_error=True,
+            max_threads=min(16, psutil.cpu_count()),
+            websocket_ping_timeout=60,
+            preventive_refresh=True,
+            memory_limit_mb=int(total_ram_gb * 1024 * 0.8)  # Use up to 80% of RAM
+        )
+    except Exception as e:
+        print(f"Critical error during startup: {e}")
+        # Attempt to clean up resources
+        _CACHE.data.clear()
+        raise