Add cached analysis results with LFS

.gitattributes

@@ -1 +1,2 @@
 *.npz filter=lfs diff=lfs merge=lfs -text
+cache/*.json filter=lfs diff=lfs merge=lfs -text

cache/cached_function_vector_results.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:332bae308b3bc0f92d10bc2389f170c88563080c0bc68de469cef8300d8e9bd7
+size 10673942

function_vectors/function_vectors_page.py

@@ -520,23 +520,157 @@ class LayerEvolutionAnalyzer:
             'layer_changes': layer_changes
         }
 
+def update_fv_cache(input_text, results):
+    cache_file = os.path.join("cache", "cached_function_vector_results.json")
+    os.makedirs("cache", exist_ok=True)
+    
+    try:
+        if os.path.exists(cache_file):
+            with open(cache_file, "r", encoding="utf-8") as f:
+                cached_data = json.load(f)
+        else:
+            cached_data = {}
+    except:
+        cached_data = {}
+            
+    # Recursive serializer to handle numpy types
+    def make_serializable(obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        if isinstance(obj, (np.float32, np.float64, np.float16)):
+            return float(obj)
+        if isinstance(obj, (np.int32, np.int64, np.int16)):
+            return int(obj)
+        if isinstance(obj, (np.bool_, bool)):
+            return bool(obj)
+        if isinstance(obj, dict):
+            return {k: make_serializable(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [make_serializable(v) for v in obj]
+        return obj
+
+    serializable_data = {
+        'attribution': {},
+        'evolution': make_serializable(results.get('evolution')),
+        'pca_explanation': results.get('pca_explanation'),
+        'evolution_explanation': results.get('evolution_explanation'),
+        'faithfulness': results.get('faithfulness', {})
+    }
+    
+    if 'attribution' in results:
+        attr = results['attribution']
+        serializable_data['attribution'] = {
+            'input_activation': make_serializable(attr.get('input_activation')),
+            'function_type_scores': make_serializable(attr.get('function_type_scores')),
+            'category_scores': make_serializable(attr.get('category_scores')),
+            'input_text': attr.get('input_text')
+        }
+        
+    cached_data[input_text] = serializable_data
+    
+    with open(cache_file, "w", encoding="utf-8") as f:
+        json.dump(cached_data, f, ensure_ascii=False, indent=4)
+    print(f"Saved FV analysis for '{input_text}' to cache.")
+
+def update_fv_cache_with_faithfulness(input_text, key, verification_results):
+    cache_file = os.path.join("cache", "cached_function_vector_results.json")
+    if not os.path.exists(cache_file): return
+    
+    # Recursive serializer to handle numpy types
+    def make_serializable(obj):
+        if isinstance(obj, np.ndarray):
+            return obj.tolist()
+        if isinstance(obj, (np.float32, np.float64, np.float16)):
+            return float(obj)
+        if isinstance(obj, (np.int32, np.int64, np.int16)):
+            return int(obj)
+        if isinstance(obj, (np.bool_, bool)):
+            return bool(obj)
+        if isinstance(obj, dict):
+            return {k: make_serializable(v) for k, v in obj.items()}
+        if isinstance(obj, list):
+            return [make_serializable(v) for v in obj]
+        return obj
+    
+    try:
+        with open(cache_file, "r", encoding="utf-8") as f:
+            cached_data = json.load(f)
+            
+        if input_text in cached_data:
+            if "faithfulness" not in cached_data[input_text]:
+                cached_data[input_text]["faithfulness"] = {}
+            
+            cached_data[input_text]["faithfulness"][key] = make_serializable(verification_results)
+            
+            with open(cache_file, "w", encoding="utf-8") as f:
+                json.dump(cached_data, f, ensure_ascii=False, indent=4)
+            print(f"Saved faithfulness for {key} to cache.")
+    except Exception as e:
+        print(f"Failed to update FV cache with faithfulness: {e}")
+
 def run_interactive_analysis(input_text, include_attribution=True, include_evolution=True, enable_ai_explanation=True):
     # A wrapper function for running the analysis from the UI.
     
+    # Check cache first
+    cache_file = os.path.join("cache", "cached_function_vector_results.json")
+    if os.path.exists(cache_file):
+        try:
+            with open(cache_file, "r", encoding="utf-8") as f:
+                cached_data = json.load(f)
+            if input_text in cached_data:
+                print(f"Loading FV analysis for '{input_text}' from cache.")
+                data = cached_data[input_text]
+                
+                results = {
+                    'evolution': data.get('evolution'),
+                    'pca_explanation': data.get('pca_explanation'),
+                    'evolution_explanation': data.get('evolution_explanation'),
+                    'faithfulness': data.get('faithfulness')
+                }
+                
+                if 'attribution' in data:
+                    attr_data = data['attribution']
+                    input_activation = np.array(attr_data['input_activation'])
+                    
+                    # Load static vectors
+                    current_lang = st.session_state.get('lang', 'en')
+                    ft_vectors, cat_vectors, error = _load_precomputed_vectors(current_lang)
+                    
+                    if not error:
+                        results['attribution'] = {
+                            'input_activation': input_activation,
+                            'function_type_scores': attr_data.get('function_type_scores'),
+                            'category_scores': attr_data.get('category_scores'),
+                            'function_types_mapping': FUNCTION_TYPES,
+                            'input_text': input_text,
+                            'category_vectors': cat_vectors,
+                            'function_type_vectors': ft_vectors
+                        }
+                
+                st.session_state.user_input_3d_data = results.get('attribution')
+                
+                # Populate faithfulness in analysis_results if needed
+                if 'faithfulness' in results and results['faithfulness']:
+                    if 'analysis_results' not in st.session_state:
+                        st.session_state.analysis_results = {}
+                    st.session_state.analysis_results['pca_faithfulness'] = results['faithfulness'].get('pca')
+                    st.session_state.analysis_results['evolution_faithfulness'] = results['faithfulness'].get('evolution')
+
+                return results
+        except Exception as e:
+            print(f"Error loading from cache: {e}")
+
     # Before running, check if models exist if not using a cached value.
-    # This check relies on the fact that caching is attempted first.
     model_path = "./models/OLMo-2-1124-7B"
     model_exists = os.path.exists(model_path)
-    # if not os.path.exists(model_path):
-    #     # We assume if the model path is missing, we are in a static environment.
-    #     # The calling function should have already checked the cache.
-    #     st.info("This live demo is running in a static environment. Only the pre-cached example prompts are available. Please select an example to view its analysis.")
-    #     return None
 
     current_lang = st.session_state.get('lang', 'en')
     
     try:
         results = _perform_analysis(input_text, include_attribution, include_evolution, current_lang, enable_ai_explanation)
+        # Save to cache
+        update_fv_cache(input_text, results)
+        
     except Exception as e:
         if not model_exists:
              st.info("This live demo is running in a static environment. Only the pre-cached example prompts are available. Please select an example to view its analysis.")
@@ -1289,6 +1423,9 @@ def display_analysis_results(results, input_text):
                             with st.spinner(tr('running_faithfulness_check_spinner')):
                                 claims = _cached_extract_fv_claims(api_config, st.session_state.explanation_part_1, "pca")
                                 verification_results = verify_fv_claims(claims, results, "pca")
+                                # Update cache
+                                if 'attribution' in results and 'input_text' in results['attribution']:
+                                     update_fv_cache_with_faithfulness(results['attribution']['input_text'], "pca", verification_results)
                         else:
                             verification_results = []
                             st.warning(tr('api_key_not_configured_warning'))
@@ -1365,6 +1502,9 @@ def display_analysis_results(results, input_text):
                             with st.spinner(tr('running_faithfulness_check_spinner')):
                                 claims = _cached_extract_fv_claims(api_config, st.session_state.explanation_part_2, "pca")
                                 verification_results = verify_fv_claims(claims, results, "pca")
+                                # Update cache
+                                if 'attribution' in results and 'input_text' in results['attribution']:
+                                     update_fv_cache_with_faithfulness(results['attribution']['input_text'], "pca", verification_results)
                         else:
                             verification_results = []
                             st.warning(tr('api_key_not_configured_warning'))
@@ -1549,6 +1689,9 @@ def display_analysis_results(results, input_text):
                                 with st.spinner(tr('running_faithfulness_check_spinner')):
                                     claims = _cached_extract_fv_claims(api_config, st.session_state.explanation_part_3, "pca")
                                     verification_results = verify_fv_claims(claims, results, "pca")
+                                # Update cache
+                                if 'attribution' in results and 'input_text' in results['attribution']:
+                                     update_fv_cache_with_faithfulness(results['attribution']['input_text'], "pca", verification_results)
                             else:
                                 verification_results = []
                                 st.warning(tr('api_key_not_configured_warning'))
@@ -1713,6 +1856,9 @@ def display_evolution_results(evolution_results):
                         with st.spinner(tr('running_faithfulness_check_spinner')):
                             claims = _cached_extract_fv_claims(api_config, st.session_state.evolution_explanation_part_1, "evolution")
                             verification_results = verify_fv_claims(claims, st.session_state.analysis_results, "evolution")
+                            # Update cache
+                            if 'attribution' in st.session_state.analysis_results and 'input_text' in st.session_state.analysis_results['attribution']:
+                                 update_fv_cache_with_faithfulness(st.session_state.analysis_results['attribution']['input_text'], "evolution", verification_results)
                     else:
                         verification_results = []
                         st.warning(tr('api_key_not_configured_warning'))
@@ -1821,6 +1967,9 @@ def display_evolution_results(evolution_results):
                         with st.spinner(tr('running_faithfulness_check_spinner')):
                             claims = _cached_extract_fv_claims(api_config, st.session_state.evolution_explanation_part_2, "evolution")
                             verification_results = verify_fv_claims(claims, st.session_state.analysis_results, "evolution")
+                            # Update cache
+                            if 'attribution' in st.session_state.analysis_results and 'input_text' in st.session_state.analysis_results['attribution']:
+                                 update_fv_cache_with_faithfulness(st.session_state.analysis_results['attribution']['input_text'], "evolution", verification_results)
                     else:
                         verification_results = []
                         st.warning(tr('api_key_not_configured_warning'))