app refactored222

app.py

@@ -114,10 +114,18 @@ def create_feature_vector(w_enc, w_dec, feature_number, weight_type, token_centr
     return feature_vector
 
 def perform_pca(_embeddings):
-    pca = PCA(n_components=1)
-    pca.fit(_embeddings.cpu().numpy())
-    pca_direction = torch.tensor(pca.components_[0], dtype=config.DTYPE, device=config.DEVICE)
-    return F.normalize(pca_direction, p=2, dim=0)
+    try:
+        pca = PCA(n_components=1)
+        embeddings_cpu = _embeddings.cpu().numpy()
+        pca.fit(embeddings_cpu)
+        pca_direction = torch.tensor(pca.components_[0], dtype=config.DTYPE, device=config.DEVICE)
+        return F.normalize(pca_direction, p=2, dim=0)
+    except Exception as e:
+        logger.error(f"Error in perform_pca: {str(e)}")
+        logger.error(traceback.format_exc())
+        raise RuntimeError(f"PCA calculation failed: {str(e)}")
+
+
 
 @torch.no_grad()
 def create_ghost_token(_feature_vector, _token_centroid, _pca_direction, target_distance, pca_weight):
@@ -174,38 +182,50 @@ def initialize_resources():
 def process_input(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp):
     global w_enc_dict, w_dec_dict
     
-    if selected_sae not in w_enc_dict or selected_sae not in w_dec_dict:
-        w_enc, w_dec = load_sae_weights(selected_sae)
-        if w_enc is None or w_dec is None:
-            return f"Failed to load SAE weights for {selected_sae}. Please try a different SAE or check your connection."
-        w_enc_dict[selected_sae] = w_enc
-        w_dec_dict[selected_sae] = w_dec
-    else:
-        w_enc, w_dec = w_enc_dict[selected_sae], w_dec_dict[selected_sae]
-    
-    token_centroid = torch.mean(token_embeddings, dim=0)
-    feature_vector = create_feature_vector(w_enc, w_dec, int(feature_number), weight_type, token_centroid, use_token_centroid, scaling_factor)
-    
-    if use_pca:
-        pca_direction = perform_pca(token_embeddings)
-        feature_vector = create_ghost_token(feature_vector, token_centroid, pca_direction, scaling_factor, pca_weight)
-    
-    closest_tokens_with_values = find_closest_tokens(
-        feature_vector, token_embeddings, tokenizer, 
-        top_k=500, num_exp=num_exp, denom_exp=denom_exp
-    )
-    
-    token_list = [token for token, _ in closest_tokens_with_values]
-    result = f"100 tokens whose embeddings produce the smallest ratio:\n\n"
-    result += f"[{', '.join(repr(token) for token in token_list[:100])}]\n\n"
-    result += "Top 500 list:\n"
-    result += "\n".join([f"{token!r}: {value:.4f}" for token, value in closest_tokens_with_values])
-    
-    return result
+    try:
+        if selected_sae not in w_enc_dict or selected_sae not in w_dec_dict:
+            w_enc, w_dec = load_sae_weights(selected_sae)
+            if w_enc is None or w_dec is None:
+                return f"Failed to load SAE weights for {selected_sae}. Please try a different SAE or check your connection."
+            w_enc_dict[selected_sae] = w_enc
+            w_dec_dict[selected_sae] = w_dec
+        else:
+            w_enc, w_dec = w_enc_dict[selected_sae], w_dec_dict[selected_sae]
+        
+        token_centroid = torch.mean(token_embeddings, dim=0)
+        feature_vector = create_feature_vector(w_enc, w_dec, int(feature_number), weight_type, token_centroid, use_token_centroid, scaling_factor)
+        
+        if use_pca:
+            logger.info("Performing PCA...")
+            try:
+                pca_direction = perform_pca(token_embeddings)
+                feature_vector = create_ghost_token(feature_vector, token_centroid, pca_direction, scaling_factor, pca_weight)
+                logger.info("PCA completed successfully.")
+            except Exception as pca_error:
+                logger.error(f"Error during PCA: {str(pca_error)}")
+                return f"Error during PCA: {str(pca_error)}"
+        
+        closest_tokens_with_values = find_closest_tokens(
+            feature_vector, token_embeddings, tokenizer, 
+            top_k=500, num_exp=num_exp, denom_exp=denom_exp
+        )
+        
+        token_list = [token for token, _ in closest_tokens_with_values]
+        result = f"100 tokens whose embeddings produce the smallest ratio:\n\n"
+        result += f"[{', '.join(repr(token) for token in token_list[:100])}]\n\n"
+        result += "Top 500 list:\n"
+        result += "\n".join([f"{token!r}: {value:.4f}" for token, value in closest_tokens_with_values])
+        
+        return result
+    except Exception as e:
+        logger.error(f"Error in process_input: {str(e)}")
+        logger.error(traceback.format_exc())
+        return f"Error: {str(e)}"
+
 
 def gradio_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("# Gemma-2B SAE Feature ExplorerX")
+        gr.Markdown("# Gemma-2B SAE Feature Explorer")
         
         with gr.Row():
             with gr.Column(scale=2):