app.py

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("Home"):
                gr.Markdown("""
                # SAErch: Sparse Autoencoder-enhanced Semantic Search

                Welcome to SAErch, an innovative approach to semantic search using Sparse Autoencoders (SAEs) trained on dense text embeddings.

                ## Key Concepts:

                1. **Sparse Autoencoders (SAEs)**: Neural networks that learn to reconstruct input data using a sparse set of features, helping to disentangle complex representations.

                2. **Feature Families**: Groups of related SAE features that represent concepts at varying levels of abstraction.

                3. **Embedding Interventions**: Technique to modify search queries by manipulating specific semantic features identified by the SAE.

                ## How It Works:

                1. SAEs are trained on embeddings from scientific paper abstracts.
                2. The SAE learns interpretable features that capture various semantic concepts.
                3. Users can interact with these features to fine-tune search queries.
                4. The system performs semantic search using the modified embeddings.

                Explore the "SAErch" tab to try out the semantic search capabilities, or dive into the "Feature Visualisation" tab to examine the learned features in more detail.

                This tool demonstrates how SAEs can bridge the gap between the semantic richness of dense embeddings and the interpretability of sparse representations, offering new possibilities for precise and explainable semantic search.
                """)


            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()