Create app.py

app.py

@@ -0,0 +1,523 @@
+import gradio as gr
+import numpy as np
+import json
+import pandas as pd
+from openai import OpenAI
+import yaml
+from typing import Optional, List, Dict, Tuple, Any
+from topk_sae import FastAutoencoder
+import torch
+import plotly.express as px
+from collections import Counter
+from huggingface_hub import hf_hub_download
+import os
+
+import os
+print(os.getenv('MODEL_REPO_ID'))
+
+# Constants
+EMBEDDING_MODEL = "text-embedding-3-small"
+d_model = 1536
+n_dirs = d_model * 6
+k = 64
+auxk = 128
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+torch.set_grad_enabled(False)
+
+# Function to download all necessary files
+def download_all_files():
+    files_to_download = [
+        "astroPH_paper_metadata.csv",
+        "csLG_feature_analysis_results_64.json",
+        "astroPH_topk_indices_64_9216_int32.npy",
+        "astroPH_64_9216.pth",
+        "astroPH_topk_values_64_9216_float16.npy",
+        "csLG_abstract_texts.json",
+        "csLG_topk_values_64_9216_float16.npy",
+        "csLG_abstract_embeddings_float16.npy",
+        "csLG_paper_metadata.csv",
+        "csLG_64_9216.pth",
+        "astroPH_abstract_texts.json",
+        "astroPH_feature_analysis_results_64.json",
+        "csLG_topk_indices_64_9216_int32.npy",
+        "astroPH_abstract_embeddings_float16.npy"
+    ]
+
+    for file in files_to_download:
+        local_path = os.path.join("data", file)
+        os.makedirs(os.path.dirname(local_path), exist_ok=True)
+        hf_hub_download(repo_id="charlieoneill/saerch-ai-data", filename=file, local_dir="data")
+        print(f"Downloaded {file}")
+
+# Load configuration and initialize OpenAI client
+download_all_files()
+# config = yaml.safe_load(open('../config.yaml', 'r'))
+# client = OpenAI(api_key=config['jwu_openai_key'])
+
+# Load the API key from the environment variable
+api_key = os.getenv('openai_key')
+
+# Ensure the API key is set
+if not api_key:
+    raise ValueError("The environment variable 'openai_key' is not set.")
+
+# Initialize the OpenAI client with the API key
+client = OpenAI(api_key=api_key)
+
+# Function to load data for a specific subject
+def load_subject_data(subject):
+    # embeddings_path = f"data/{subject}_abstract_embeddings.npy"
+    # texts_path = f"data/{subject}_abstract_texts.json"
+    # feature_analysis_path = f"data/{subject}_feature_analysis_results_{k}.json"
+    # metadata_path = f'data/{subject}_paper_metadata.csv'
+    # topk_indices_path = f"data/{subject}_topk_indices_{k}_{n_dirs}.npy"
+    # topk_values_path = f"data/{subject}_topk_values_{k}_{n_dirs}.npy"
+
+    embeddings_path = f"data/{subject}_abstract_embeddings_float16.npy"
+    texts_path = f"data/{subject}_abstract_texts.json"
+    feature_analysis_path = f"data/{subject}_feature_analysis_results_{k}.json"
+    metadata_path = f'data/{subject}_paper_metadata.csv'
+    topk_indices_path = f"data/{subject}_topk_indices_{k}_{n_dirs}_int32.npy"
+    topk_values_path = f"data/{subject}_topk_values_{k}_{n_dirs}_float16.npy"
+
+    # abstract_embeddings = np.load(embeddings_path)
+    # with open(texts_path, 'r') as f:
+    #     abstract_texts = json.load(f)
+    # with open(feature_analysis_path, 'r') as f:
+    #     feature_analysis = json.load(f)
+    # df_metadata = pd.read_csv(metadata_path)
+    # topk_indices = np.load(topk_indices_path)
+    # topk_values = np.load(topk_values_path)
+
+    abstract_embeddings = np.load(embeddings_path).astype(np.float32)  # Load float16 and convert to float32
+    with open(texts_path, 'r') as f:
+        abstract_texts = json.load(f)
+    with open(feature_analysis_path, 'r') as f:
+        feature_analysis = json.load(f)
+    df_metadata = pd.read_csv(metadata_path)
+    topk_indices = np.load(topk_indices_path)  # Already in int32, no conversion needed
+    topk_values = np.load(topk_values_path).astype(np.float32)
+
+    model_filename = f"{subject}_64_9216.pth"
+    model_path = os.path.join("data", model_filename)
+
+    ae = FastAutoencoder(n_dirs, d_model, k, auxk, multik=0).to(device)
+    ae.load_state_dict(torch.load(model_path))
+    ae.eval()
+
+    weights = torch.load(model_path)
+    decoder = weights['decoder.weight'].cpu().numpy()
+    del weights
+
+    return {
+        'abstract_embeddings': abstract_embeddings,
+        'abstract_texts': abstract_texts,
+        'feature_analysis': feature_analysis,
+        'df_metadata': df_metadata,
+        'topk_indices': topk_indices,
+        'topk_values': topk_values,
+        'ae': ae,
+        'decoder': decoder
+    }
+
+# Load data for both subjects
+subject_data = {
+    'astroPH': load_subject_data('astroPH'),
+    'csLG': load_subject_data('csLG')
+}
+
+# Update existing functions to use the selected subject's data
+def get_embedding(text: Optional[str], model: str = EMBEDDING_MODEL) -> Optional[np.ndarray]:
+    try:
+        embedding = client.embeddings.create(input=[text], model=model).data[0].embedding
+        return np.array(embedding, dtype=np.float32)
+    except Exception as e:
+        print(f"Error getting embedding: {e}")
+        return None
+
+def intervened_hidden_to_intervened_embedding(topk_indices, topk_values, ae):
+    with torch.no_grad():
+        return ae.decode_sparse(topk_indices, topk_values)
+
+# Function definitions for feature activation, co-occurrence, styling, etc.
+def get_feature_activations(subject, feature_index, m=5, min_length=100):
+    abstract_texts = subject_data[subject]['abstract_texts']
+    abstract_embeddings = subject_data[subject]['abstract_embeddings']
+    topk_indices = subject_data[subject]['topk_indices']
+    topk_values = subject_data[subject]['topk_values']
+
+    doc_ids = abstract_texts['doc_ids']
+    abstracts = abstract_texts['abstracts']
+    
+    feature_mask = topk_indices == feature_index
+    activated_indices = np.where(feature_mask.any(axis=1))[0]
+    activation_values = np.where(feature_mask, topk_values, 0).max(axis=1)
+    
+    sorted_activated_indices = activated_indices[np.argsort(-activation_values[activated_indices])]
+    
+    top_m_abstracts = []
+    top_m_indices = []
+    for i in sorted_activated_indices:
+        if len(abstracts[i]) > min_length:
+            top_m_abstracts.append((doc_ids[i], abstracts[i], activation_values[i]))
+            top_m_indices.append(i)
+        if len(top_m_abstracts) == m:
+            break
+    
+    return top_m_abstracts
+
+def calculate_co_occurrences(subject, target_index, n_features=9216):
+    topk_indices = subject_data[subject]['topk_indices']
+
+    mask = np.any(topk_indices == target_index, axis=1)
+    co_occurring_indices = topk_indices[mask].flatten()
+    co_occurrences = Counter(co_occurring_indices)
+    del co_occurrences[target_index]
+    result = np.zeros(n_features, dtype=int)
+    result[list(co_occurrences.keys())] = list(co_occurrences.values())
+    return result
+
+def style_dataframe(df: pd.DataFrame, is_top: bool) -> pd.DataFrame:
+    cosine_values = df['Cosine similarity'].astype(float)
+    min_val = cosine_values.min()
+    max_val = cosine_values.max()
+    
+    def color_similarity(val):
+        val = float(val)
+        # Normalize the value between 0 and 1
+        if is_top:
+            normalized_val = (val - min_val) / (max_val - min_val)
+        else:
+            # For bottom correlated, reverse the normalization
+            normalized_val = (max_val - val) / (max_val - min_val)
+        
+        # Adjust the color intensity to avoid zero intensity
+        color_intensity = 0.2 + (normalized_val * 0.8)  # This ensures the range is from 0.2 to 1.0
+        
+        if is_top:
+            color = f'background-color: rgba(0, 255, 0, {color_intensity:.2f})'
+        else:
+            color = f'background-color: rgba(255, 0, 0, {color_intensity:.2f})'
+        return color
+
+    return df.style.applymap(color_similarity, subset=['Cosine similarity'])
+
+def get_feature_from_index(subject, index):
+    feature = next((f for f in subject_data[subject]['feature_analysis'] if f['index'] == index), None)
+    return feature
+
+def visualize_feature(subject, index):
+    feature = next((f for f in subject_data[subject]['feature_analysis'] if f['index'] == index), None)
+    if feature is None:
+        return "Invalid feature index", None, None, None, None, None, None
+
+    output = f"# {feature['label']}\n\n"
+    output += f"* Pearson correlation: {feature['pearson_correlation']:.4f}\n\n"
+    output += f"* Density: {feature['density']:.4f}\n\n"
+
+    # Top m abstracts
+    top_m_abstracts = get_feature_activations(subject, index)
+    
+    # Create dataframe for top abstracts
+    df_data = [
+        {"Title": m[1].split('\n\n')[0], "Activation value": f"{m[2]:.4f}"}
+        for m in top_m_abstracts
+    ]
+    df_top_abstracts = pd.DataFrame(df_data)
+
+    # Activation value distribution
+    topk_indices = subject_data[subject]['topk_indices']
+    topk_values = subject_data[subject]['topk_values']
+
+    activation_values = np.where(topk_indices == index, topk_values, 0).max(axis=1)
+    fig2 = px.histogram(x=activation_values, nbins=50)
+    fig2.update_layout(
+        #title=f'{feature["label"]}',
+        xaxis_title='Activation value',
+        yaxis_title=None,
+        yaxis_type='log',
+        height=220,
+    )
+
+    # Correlated features
+    decoder = subject_data[subject]['decoder']
+    feature_vector = decoder[:, index]
+    decoder_without_feature = np.delete(decoder, index, axis=1)
+    cosine_similarities = np.dot(feature_vector, decoder_without_feature) / (np.linalg.norm(decoder_without_feature, axis=0) * np.linalg.norm(feature_vector))
+
+    topk = 5
+    topk_indices_cosine = np.argsort(-cosine_similarities)[:topk]
+    topk_values_cosine = cosine_similarities[topk_indices_cosine]
+
+    # Create dataframe for top 5 correlated features
+    df_top_correlated = pd.DataFrame({
+        "Feature": [get_feature_from_index(subject, i)['label'] for i in topk_indices_cosine],
+        "Cosine similarity": [f"{v:.4f}" for v in topk_values_cosine]
+    })
+    df_top_correlated_styled = style_dataframe(df_top_correlated, is_top=True)
+
+    bottomk = 5
+    bottomk_indices_cosine = np.argsort(cosine_similarities)[:bottomk]
+    bottomk_values_cosine = cosine_similarities[bottomk_indices_cosine]
+
+    # Create dataframe for bottom 5 correlated features
+    df_bottom_correlated = pd.DataFrame({
+        "Feature": [get_feature_from_index(subject, i)['label'] for i in bottomk_indices_cosine],
+        "Cosine similarity": [f"{v:.4f}" for v in bottomk_values_cosine]
+    })
+    df_bottom_correlated_styled = style_dataframe(df_bottom_correlated, is_top=False)
+
+    # Co-occurrences
+    co_occurrences = calculate_co_occurrences(subject, index)
+    topk = 5
+    topk_indices_co_occurrence = np.argsort(-co_occurrences)[:topk]
+    topk_values_co_occurrence = co_occurrences[topk_indices_co_occurrence]
+
+    # Create dataframe for top 5 co-occurring features
+    df_co_occurrences = pd.DataFrame({
+        "Feature": [get_feature_from_index(subject, i)['label'] for i in topk_indices_co_occurrence],
+        "Co-occurrences": topk_values_co_occurrence
+    })
+
+    return output, df_top_abstracts, df_top_correlated_styled, df_bottom_correlated_styled, df_co_occurrences, fig2
+
+# Modify the main interface function
+def create_interface():
+    custom_css = """
+    #custom-slider-* {
+        background-color: #ffe6e6;
+    }
+    """
+
+    with gr.Blocks(css=custom_css) as demo:
+        subject = gr.Dropdown(choices=['astroPH', 'csLG'], label="Select Subject", value='astroPH')
+        
+        with gr.Tabs():
+            with gr.Tab("SAErch"):
+                input_text = gr.Textbox(label="input")
+                search_results_state = gr.State([])
+                feature_values_state = gr.State([])
+                feature_indices_state = gr.State([])
+                manually_added_features_state = gr.State([])
+
+                def update_search_results(feature_values, feature_indices, manually_added_features, current_subject):
+                    ae = subject_data[current_subject]['ae']
+                    abstract_embeddings = subject_data[current_subject]['abstract_embeddings']
+                    abstract_texts = subject_data[current_subject]['abstract_texts']
+                    df_metadata = subject_data[current_subject]['df_metadata']
+
+                    # Combine manually added features with query-generated features
+                    all_indices = []
+                    all_values = []
+                    
+                    # Add manually added features first
+                    for index in manually_added_features:
+                        if index not in all_indices:
+                            all_indices.append(index)
+                            all_values.append(feature_values[feature_indices.index(index)] if index in feature_indices else 0.0)
+                    
+                    # Add remaining query-generated features
+                    for index, value in zip(feature_indices, feature_values):
+                        if index not in all_indices:
+                            all_indices.append(index)
+                            all_values.append(value)
+
+                    # Reconstruct query embedding
+                    topk_indices = torch.tensor(all_indices).to(device)
+                    topk_values = torch.tensor(all_values).to(device)
+                    
+                    intervened_embedding = intervened_hidden_to_intervened_embedding(topk_indices, topk_values, ae)
+                    intervened_embedding = intervened_embedding.cpu().numpy().flatten()
+
+                    # Perform similarity search
+                    sims = np.dot(abstract_embeddings, intervened_embedding)
+                    topk_indices_search = np.argsort(sims)[::-1][:10]
+                    doc_ids = abstract_texts['doc_ids']
+                    topk_doc_ids = [doc_ids[i] for i in topk_indices_search]
+
+                    # Prepare search results
+                    search_results = []
+                    for doc_id in topk_doc_ids:
+                        metadata = df_metadata[df_metadata['arxiv_id'] == doc_id].iloc[0]
+                        title = metadata['title'].replace('[', '').replace(']', '')
+                        search_results.append([
+                            title,
+                            int(metadata['citation_count']),
+                            int(metadata['year'])
+                        ])
+
+                    return search_results, all_values, all_indices
+
+                @gr.render(inputs=[input_text, search_results_state, feature_values_state, feature_indices_state, manually_added_features_state, subject])
+                def show_components(text, search_results, feature_values, feature_indices, manually_added_features, current_subject):
+                    if len(text) == 0:
+                        return gr.Markdown("## No Input Provided")
+
+                    if not search_results or text != getattr(show_components, 'last_query', None):
+                        show_components.last_query = text
+                        query_embedding = get_embedding(text)
+
+                        ae = subject_data[current_subject]['ae']
+                        with torch.no_grad():
+                            recons, z_dict = ae(torch.tensor(query_embedding).unsqueeze(0).to(device))
+                            topk_indices = z_dict['topk_indices'][0].cpu().numpy()
+                            topk_values = z_dict['topk_values'][0].cpu().numpy()
+
+                        feature_values = topk_values.tolist()
+                        feature_indices = topk_indices.tolist()
+                        search_results, feature_values, feature_indices = update_search_results(feature_values, feature_indices, manually_added_features, current_subject)
+
+                    with gr.Row():
+                        with gr.Column(scale=2):
+                            df = gr.Dataframe(
+                                headers=["Title", "Citation Count", "Year"],
+                                value=search_results,
+                                label="Top 10 Search Results"
+                            )
+
+                            feature_search = gr.Textbox(label="Search Feature Labels")
+                            feature_matches = gr.CheckboxGroup(label="Matching Features", choices=[])
+                            add_button = gr.Button("Add Selected Features")
+
+                            def search_feature_labels(search_text):
+                                if not search_text:
+                                    return gr.CheckboxGroup(choices=[])
+                                matches = [f"{f['label']} ({f['index']})" for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
+                                return gr.CheckboxGroup(choices=matches[:10])
+
+                            feature_search.change(search_feature_labels, inputs=[feature_search], outputs=[feature_matches])
+
+                            def on_add_features(selected_features, current_values, current_indices, manually_added_features):
+                                if selected_features:
+                                    new_indices = [int(f.split('(')[-1].strip(')')) for f in selected_features]
+                                    
+                                    # Add new indices to manually_added_features if they're not already there
+                                    manually_added_features = list(dict.fromkeys(manually_added_features + new_indices))
+                                    
+                                    return gr.CheckboxGroup(value=[]), current_values, current_indices, manually_added_features
+                                return gr.CheckboxGroup(value=[]), current_values, current_indices, manually_added_features
+
+                            add_button.click(
+                                on_add_features,
+                                inputs=[feature_matches, feature_values_state, feature_indices_state, manually_added_features_state],
+                                outputs=[feature_matches, feature_values_state, feature_indices_state, manually_added_features_state]
+                            )
+
+                        with gr.Column(scale=1):
+                            update_button = gr.Button("Update Results")
+                            sliders = []
+                            for i, (value, index) in enumerate(zip(feature_values, feature_indices)):
+                                feature = next((f for f in subject_data[current_subject]['feature_analysis'] if f['index'] == index), None)
+                                label = f"{feature['label']} ({index})" if feature else f"Feature {index}"
+                                
+                                # Add prefix and change color for manually added features
+                                if index in manually_added_features:
+                                    label = f"[Custom] {label}"
+                                    slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=value, label=label, key=f"slider-{index}", elem_id=f"custom-slider-{index}")
+                                else:
+                                    slider = gr.Slider(minimum=0, maximum=1, step=0.01, value=value, label=label, key=f"slider-{index}")
+                                
+                                sliders.append(slider)
+
+                    def on_slider_change(*values):
+                        manually_added_features = values[-1]
+                        slider_values = list(values[:-1])
+                        
+                        # Reconstruct feature_indices based on the order of sliders
+                        reconstructed_indices = [int(slider.label.split('(')[-1].split(')')[0]) for slider in sliders]
+                        
+                        new_results, new_values, new_indices = update_search_results(slider_values, reconstructed_indices, manually_added_features, current_subject)
+                        return new_results, new_values, new_indices, manually_added_features
+
+                    update_button.click(
+                        on_slider_change,
+                        inputs=sliders + [manually_added_features_state],
+                        outputs=[search_results_state, feature_values_state, feature_indices_state, manually_added_features_state]
+                    )
+
+                    return [df, feature_search, feature_matches, add_button, update_button] + sliders
+
+            with gr.Tab("Feature Visualisation"):
+                gr.Markdown("# Feature Visualiser")
+                with gr.Row():
+                    feature_search = gr.Textbox(label="Search Feature Labels")
+                    feature_matches = gr.CheckboxGroup(label="Matching Features", choices=[])
+                    visualize_button = gr.Button("Visualize Feature")
+                
+                feature_info = gr.Markdown()
+                abstracts_heading = gr.Markdown("## Top 5 Abstracts")
+                top_abstracts = gr.Dataframe(
+                    headers=["Title", "Activation value"],
+                    interactive=False
+                )
+                
+                gr.Markdown("## Correlated Features")
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Top 5 Correlated Features")
+                        top_correlated = gr.Dataframe(
+                            headers=["Feature", "Cosine similarity"],
+                            interactive=False
+                        )
+                    with gr.Column(scale=1):
+                        gr.Markdown("### Bottom 5 Correlated Features")
+                        bottom_correlated = gr.Dataframe(
+                            headers=["Feature", "Cosine similarity"],
+                            interactive=False
+                        )
+                
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        gr.Markdown("## Top 5 Co-occurring Features")
+                        co_occurring_features = gr.Dataframe(
+                            headers=["Feature", "Co-occurrences"],
+                            interactive=False
+                        )
+                    with gr.Column(scale=1):
+                        gr.Markdown(f"## Activation Value Distribution")
+                        activation_dist = gr.Plot()
+
+                def search_feature_labels(search_text, current_subject):
+                    if not search_text:
+                        return gr.CheckboxGroup(choices=[])
+                    matches = [f"{f['label']} ({f['index']})" for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
+                    return gr.CheckboxGroup(choices=matches[:10])
+
+                feature_search.change(search_feature_labels, inputs=[feature_search, subject], outputs=[feature_matches])
+
+                def on_visualize(selected_features, current_subject):
+                    if not selected_features:
+                        return "Please select a feature to visualize.", None, None, None, None, None, "", []
+                    
+                    # Extract the feature index from the selected feature string
+                    feature_index = int(selected_features[0].split('(')[-1].strip(')'))
+                    feature_info, top_abstracts, top_correlated, bottom_correlated, co_occurring_features, activation_dist = visualize_feature(current_subject, feature_index)
+                    
+                    # Return the visualization results along with empty values for search box and checkbox
+                    return feature_info, top_abstracts, top_correlated, bottom_correlated, co_occurring_features, activation_dist, "", []
+
+                visualize_button.click(
+                    on_visualize,
+                    inputs=[feature_matches, subject],
+                    outputs=[feature_info, top_abstracts, top_correlated, bottom_correlated, co_occurring_features, activation_dist, feature_search, feature_matches]
+                )
+
+        # Add logic to update components when subject changes
+        def on_subject_change(new_subject):
+            # Clear all states and return empty values for all components
+            return [], [], [], [], "", [], "", [], None, None, None, None, None, None
+
+        subject.change(
+            on_subject_change,
+            inputs=[subject],
+            outputs=[search_results_state, feature_values_state, feature_indices_state, manually_added_features_state, 
+                     input_text, feature_matches, feature_search, feature_matches, 
+                     feature_info, top_abstracts, top_correlated, bottom_correlated, co_occurring_features, activation_dist]
+        )
+
+    return demo
+
+# Launch the interface
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()