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 networkx as nx
# import plotly.graph_objs as go
# from ast import literal_eval as make_tuple
# import random

# 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",
# # "csLG_clean_families_64_9216.json",
# # "astroPH_clean_families_64_9216.json",
# # "astroPH_family_analysis_64_9216.json",
# "csLG_family_analysis_64_9216.json"
# ]

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

# # 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_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"
# norms_path = f"data/{subject}_norms_{k}_{n_dirs}.npy"
# topk_values_path = f"data/{subject}_topk_values_{k}_{n_dirs}_float16.npy"
# families_path = f"data/{subject}_clean_families_{k}_{n_dirs}.json"
# family_analysis_path = f"data/{subject}_family_analysis_{k}_{n_dirs}.json"
# nns_32to64 = json.load(open(f"data/{subject}_nns_32to64.json"))
# nns_16to32 = json.load(open(f"data/{subject}_nns_16to32.json"))
# nns_16to64 = json.load(open(f"data/{subject}_nns_16to64.json"))

# 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)
# norms = np.load(norms_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

# with open(family_analysis_path, 'r') as f:
# family_analysis = json.load(f)


# 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,
# 'norms': norms,
# 'nns_32to64': nns_32to64,
# 'nns_16to64': nns_16to64,
# 'ae': ae,
# 'decoder': decoder,
# # 'feature_families': feature_families,
# 'family_analysis': family_analysis
# }

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 networkx as nx
import plotly.graph_objs as go
from ast import literal_eval as make_tuple
import random
import math

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",
 # "csLG_clean_families_64_9216.json",
 # "astroPH_clean_families_64_9216.json",
 "astroPH_family_analysis_64_9216.json",
 "csLG_family_analysis_64_9216.json",
 "csLG_nns_32to64.json",
 "csLG_nns_16to32.json",
 "csLG_nns_16to64.json",
 "astroPH_nns_32to64.json",
 "astroPH_nns_16to32.json",
 "astroPH_nns_16to64.json",
 "csLG_norms_64_9216_float16.npy",
 "astroPH_norms_64_9216_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()

# 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_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"
 norms_path = f"data/{subject}_norms_{k}_{n_dirs}_float16.npy"
 topk_values_path = f"data/{subject}_topk_values_{k}_{n_dirs}_float16.npy"
 families_path = f"data/{subject}_clean_families_{k}_{n_dirs}.json"
 family_analysis_path = f"data/{subject}_family_analysis_{k}_{n_dirs}.json"
 nns_32to64 = json.load(open(f"data/{subject}_nns_32to64.json"))
 nns_16to32 = json.load(open(f"data/{subject}_nns_16to32.json"))
 nns_16to64 = json.load(open(f"data/{subject}_nns_16to64.json"))

 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)
 norms = np.load(norms_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

 with open(family_analysis_path, 'r') as f:
 family_analysis = json.load(f)


 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,
 'norms': norms,
 'nns_32to64': nns_32to64,
 'nns_16to64': nns_16to64,
 'ae': ae,
 'decoder': decoder,
 # 'feature_families': feature_families,
 'family_analysis': family_analysis
 }

# 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']
 norms = subject_data[subject]['norms']

 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=np.float32)
 result[list(co_occurrences.keys())] = list(co_occurrences.values())
 result[list(co_occurrences.keys())] /= np.minimum(norms[list(co_occurrences.keys())], norms[target_index]) 
 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 with clickable links
 df_data = []
 for doc_id, abstract, activation_value in top_m_abstracts:
 title = abstract.split('\n\n')[0]
 title = title.replace('[', '').replace(']', '')
 title = title.replace("'", "")
 title = title.replace('"', '')
 url_id = doc_id.replace('_arXiv.txt', '')
 if 'astro-ph' in url_id:
 url_id = url_id.split('astro-ph')[1]
 url = f"https://arxiv.org/abs/astro-ph/{url_id}"
 else:
 if '.' in doc_id:
 url = f"https://arxiv.org/abs/{url_id}"
 else:
 url = f"https://arxiv.org/abs/hep-ph/{url_id}"
 
 linked_title = f"[{title}]({url})"
 df_data.append({"Title": linked_title, "Activation value": activation_value})
 
 df_top_abstracts = pd.DataFrame(df_data)
 styled_top_abstracts = df_top_abstracts.style.format({
 "Activation value": "{:.4f}"
 })

 # 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]

 bottomk = 5
 bottomk_indices_cosine = np.argsort(cosine_similarities)[:bottomk]
 bottomk_values_cosine = cosine_similarities[bottomk_indices_cosine]
 
 df_top_correlated = pd.DataFrame({
 "Feature": [get_feature_from_index(subject, i)['label'] for i in topk_indices_cosine],
 "Cosine similarity": topk_values_cosine
 })
 df_top_correlated_styled = style_dataframe(df_top_correlated, is_top=True)

 # 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": 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
 })
 df_co_occurrences_styled = df_co_occurrences.style.format({
 "Co-occurrences": "{:.2f}" # 2 decimal points
 })

 # Add new code for feature splitting
 nns_16to64 = subject_data[subject]['nns_16to64']
 nns_32to64 = subject_data[subject]['nns_32to64']

 # Get nearest neighbors for 16 and 32
 #nn_16 = nns_16to64[str(index)]

 # this is really involved it's a lot easier the other direction
 nn_16 = []
 for key in nns_16to64.keys():
 for match in nns_16to64[key]:
 if index == match['feature'][0]:
 nn_16.append([key, float(match['similarity'])])
 
 #nn_32 = nns_32to64[str(index)]
 nn_32 = []
 for key in nns_32to64.keys():
 for match in nns_32to64[key]:
 if index == match['feature'][0]:
 nn_32.append([key, float(match['similarity'])])

 # Create dataframes for 16 and 32 nearest neighbors
 try:
 df_16 = pd.DataFrame(nn_16, columns=["Feature", "Cosine Similarity"])
 df_16 = df_16.style.format({"Cosine Similarity": "{:.4f}"})
 except:
 df_16 = pd.DataFrame(["No Match"], columns=["Feature"])

 try:
 df_32 = pd.DataFrame(nn_32, columns=["Feature", "Cosine Similarity"])
 df_32 = df_32.style.format({"Cosine Similarity": "{:.4f}"})
 except:
 df_32 = pd.DataFrame(["No Match"], columns=["Feature"])

 return output, styled_top_abstracts, df_top_correlated_styled, df_bottom_correlated_styled, df_co_occurrences_styled, fig2, df_16, df_32

def create_interactive_directed_graph(family):
 matrix = np.array(family['matrix'])
 matrix[matrix < 0.07] = 0
 densities = family['densities']
 for i in range(len(densities)):
 for j in range(len(densities)):
 if densities[i] < densities[j]:
 matrix[i][j] = 0
 
 G = nx.from_numpy_array(matrix, create_using=nx.DiGraph())

 num_nodes = len(family['feature_f1'])
 all_f1s = family['feature_pearson'] + [family['family_pearson']]
 node_info = {i: {"name": f"{family['feature_names'][i]}", "density": family['densities'][i], "pearson": all_f1s[i]} for i in range(num_nodes)}
 nx.set_node_attributes(G, node_info)

 # Create node trace
 node_x = []
 node_y = []
 node_text = []
 node_size = []
 node_color = []
 pos = nx.spring_layout(G, k = np.sqrt(1/num_nodes) * 3)
 for node in G.nodes():
 x, y = pos[node]
 node_x.append(x)
 node_y.append(y)
 node_text.append(G.nodes[node]['name'] + "<br>log density: " + str(round(np.log10(G.nodes[node]['density'] + 1e-5), 3)))
 node_size.append((np.log10(G.nodes[node]['density'] + 1e-5) + 6) * 10)
 node_color.append(G.nodes[node]['pearson'])

 node_trace = go.Scatter(
 x=node_x, y=node_y,
 mode='markers',
 hoverinfo='text',
 marker=dict(
 showscale=True,
 colorscale='purples',
 size=node_size, # Set node marker size to node['f1']
 color=node_color,
 cmin = 0,
 cmax = 1,
 colorbar=dict(
 thickness=15,
 title='Pearson Correlation',
 xanchor='left',
 titleside='right',
 ),
 line_width=2,
 opacity = 1,),
 opacity = 1)

 node_trace.text = node_text

 # Create edge trace
 edge_traces = []
 annotations = []
 for edge in G.edges():
 x0, y0 = pos[edge[0]]
 x1, y1 = pos[edge[1]]
 weight = matrix[edge[0], edge[1]]

 # Calculate offset (adjust this value to move arrows further from or closer to nodes)
 offset = 0.00
 start_x = x0
 start_y = y0
 end_x = x1
 end_y = y1

 # # Calculate new start and end points
 # if start_x > end_x:
 # start_x = x0 - offset
 # end_x = x0 + offset
 # else:
 # start_x = x0 + offset
 # end_x = x1 - offset
 # if start_y > end_y:
 # start_y = y0 - offset
 # end_y = y1 + offset
 # else:
 # start_y = y0 + offset
 # end_y = y1 - offset

 edge_trace = go.Scatter(
 x=[start_x, end_x, None],
 y=[start_y, end_y, None],
 line=dict(width=weight * 20, color='#888'), # Multiply weight by 20 for better visibility
 hovertext="weight: " + str(round(weight, 3)), # Set the hover text to the edge weight
 mode='lines',
 line_shape='spline',
 opacity = 0.5,
 )
 edge_traces.append(edge_trace)
 
 annotation = dict(
 ax=start_x,
 ay=start_y,
 x=end_x,
 y=end_y,
 xref='x',
 yref='y',
 axref='x',
 ayref='y',
 showarrow=True,
 arrowhead=4,
 arrowsize=4, #max(min(weight * 3, 0.3), 2), # Reduced from 30 to 10
 arrowwidth=1, # Reduced from 30 to 2
 arrowcolor='#999',
 opacity = 1,
 )
 annotations.append(annotation)

 annotation_trace = go.Scatter(x=[], y=[], mode='markers', hoverinfo='none', marker=dict(opacity=0))

 # Create the figure
 fig = go.Figure(data=[annotation_trace, *edge_traces, node_trace],
 layout=go.Layout(
 showlegend=False,
 hovermode='closest',
 margin=dict(b=20,l=5,r=5,t=40),
 xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
 yaxis=dict(showgrid=False, zeroline=False, showticklabels=False)),
 )
 fig.update_xaxes(showline=False, linewidth=0, gridcolor='white')
 fig.update_yaxes(showline=False, linewidth=0, gridcolor='white')
 fig.update_layout(
 plot_bgcolor='white',
 annotations=annotations,
 )
 return fig

# 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. This tool builds upon recent advancements in the application of SAEs to language models and 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. SAEs have shown promising results in uncovering interpretable features in language models.

 2. **Feature Families**: Groups of related SAE features that represent concepts at varying levels of abstraction, allowing for multi-scale semantic analysis and manipulation.

 3. **Embedding Interventions**: Technique to modify search queries by manipulating specific semantic features identified by the SAE, enabling fine-grained control over query semantics.

 ## How It Works:

 1. SAEs are trained on embeddings from scientific paper abstracts, learning interpretable features that capture various semantic concepts.
 2. Users can interact with these features to fine-tune search queries.
 3. The system performs semantic search using the modified embeddings, allowing for more precise and controllable results.

 ## Key References:

 - [Towards Monosemanticity: Decomposing Language Models With Dictionary Learning](https://transformer-circuits.pub/2023/monosemantic-features) - Anthropic's pioneering work on applying SAEs to language models.
 - [Prism: Mapping Interpretable Concepts and Features in a Latent Space of Language](https://thesephist.com/posts/prism/#caveats-and-limitations) - An early application of SAEs to embeddings, demonstrating their potential for interpretable concept mapping.
 - [Scaling and Evaluating Sparse Autoencoders](https://arxiv.org/html/2406.04093v1) - OpenAI's research on scaling SAEs, showcasing the effectiveness of top-k SAEs.

 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(']', '')
 title = title.replace("'", "")

 url_id = doc_id.replace('_arXiv.txt', '')
 if 'astro-ph' in url_id:
 url_id = url_id.split('astro-ph')[1]
 url = f"https://arxiv.org/abs/astro-ph/{url_id}"
 else:
 if '.' in doc_id:
 url = f"https://arxiv.org/abs/{doc_id.replace('_arXiv.txt', '')}"
 else:
 url = f"https://arxiv.org/abs/hep-ph/{doc_id.replace('_arXiv.txt', '')}"
 
 linked_title = f"[{title}]({url})"
 
 search_results.append([
 linked_title,
 int(metadata['citation_count']),
 int(metadata['year'])
 ])

 # Convert search_results to a DataFrame and apply styling
 df_search_results = pd.DataFrame(search_results, columns=["Title", "Citation Count", "Year"])
 styled_search_results = df_search_results.style.format({
 "Citation Count": "{:.0f}", # Keep as integer
 "Year": "{:.0f}" # Keep as integer
 })

 return styled_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",
 datatype=["markdown", "number", "number"], # Add this line
 wrap=True # Add this line to ensure long titles don't get cut off
 )

 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 for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
 # matches = sorted(matches, key=lambda x: x['pearson_correlation'], reverse=True)
 # matches = [f"{f['label']} ({f['index']})" for f in matches]
 
 # return gr.CheckboxGroup(choices=matches[:10])

 def search_feature_labels(search_text):
 if not search_text:
 return gr.CheckboxGroup(choices=[])
 matches = [f for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
 for match in matches:
 if math.isnan(match['pearson_correlation']):
 match['pearson_correlation'] = 0
 matches = sorted(matches, key=lambda x: x['pearson_correlation'], reverse=True)
 matches = [f"{f['label']} ({f['index']})" for f in matches]
 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}"
 
 # Transform the value to a 0-1 range
 transformed_value = max(0.01, min(1, value)) # Ensure value is between 0.01 and 1
 linear_value = (np.log10(transformed_value) + 2) / 2 # Map 0.01-1 to 0-1
 
 # 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=linear_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=linear_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])
 
 # Transform slider values back to original scale
 transformed_values = [10 ** ((2 * float(v)) - 2) for v in slider_values]
 
 # 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(transformed_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.Tabs():
 with gr.Tab("Individual Features"):
 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"],
 datatype=["markdown", "number"],
 interactive=False,
 wrap=True
 )
 
 gr.Markdown("## Feature Splitting")
 with gr.Row():
 with gr.Column(scale=1):
 gr.Markdown("### Best Match in SAE16")
 nn_16_table = gr.Dataframe(
 headers=["Feature", "Cosine Similarity"],
 interactive=False
 )
 with gr.Column(scale=1):
 gr.Markdown("### Best Match in SAE32")
 nn_32_table = gr.Dataframe(
 headers=["Feature", "Cosine Similarity"],
 interactive=False
 )

 with gr.Row():
 with gr.Column(scale=1):
 gr.Markdown("## Top 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()

 gr.Markdown("## Correlated Features")
 with gr.Row():
 with gr.Column(scale=1):
 gr.Markdown("### Top Correlated Features")
 top_correlated = gr.Dataframe(
 headers=["Feature", "Cosine similarity"],
 interactive=False
 )
 with gr.Column(scale=1):
 gr.Markdown("### Bottom Correlated Features")
 bottom_correlated = gr.Dataframe(
 headers=["Feature", "Cosine similarity"],
 interactive=False
 )
 
 
 

 # def search_feature_labels(search_text, current_subject):
 # if not search_text:
 # return gr.CheckboxGroup(choices=[])
 # matches = [f for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
 # matches = sorted(matches, key=lambda x: x['pearson_correlation'], reverse=True)
 # matches = [f"{f['label']} ({f['index']})" for f in matches]
 
 # return gr.CheckboxGroup(choices=matches[:10])

 def search_feature_labels(search_text, current_subject):
 if not search_text:
 return gr.CheckboxGroup(choices=[])
 matches = [f for f in subject_data[current_subject]['feature_analysis'] if search_text.lower() in f['label'].lower()]
 for match in matches:
 if math.isnan(match['pearson_correation']):
 match['pearson_correlation'] = 0
 matches = sorted(matches, key=lambda x: x['pearson_correlation'], reverse=True)
 matches = [f"{f['label']} ({f['index']})" for f in matches]
 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, nn_16, nn_32 = 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, "", [], nn_16, nn_32

 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, nn_16_table, nn_32_table]
 )

 with gr.Tab("Feature Families"):
 gr.Markdown("# Feature Families")
 
 with gr.Row():
 family_search = gr.Textbox(label="Search Feature Families")
 family_matches = gr.CheckboxGroup(label="Matching Feature Families", choices=[])
 visualize_family_button = gr.Button("Visualize Feature Family")
 
 
 family_dataframe = gr.Dataframe(
 headers=["Feature", "Parent Co-Occurrence", "F1 Score", "Pearson"],
 datatype=["markdown", "number", "number", "number"],
 label="Family and Child Features"
 )

 gr.Markdown("# Family Graph")
 graph_plot = gr.Plot(label="Directed Graph")

 # family_info = gr.Markdown()

 # def search_feature_families(search_text, current_subject):
 # family_analysis = subject_data[current_subject]['family_analysis']
 # if not search_text:
 # return gr.CheckboxGroup(choices=[])
 # matches = [family for family in family_analysis if search_text.lower() in family['superfeature'].lower()]
 
 # matches = sorted(matches, key=lambda x: x['family_pearson'], reverse=True)
 # matches = [family['superfeature'] for family in matches]
 # matches = list(dict.fromkeys(matches))
 # return gr.CheckboxGroup(choices=matches[:10]) # Limit to top 10 matches

 def search_feature_families(search_text, current_subject):
 family_analysis = subject_data[current_subject]['family_analysis']
 if not search_text:
 return gr.CheckboxGroup(choices=[])
 matches = [family for family in family_analysis if search_text.lower() in family['superfeature'].lower()]
 for family in matches:
 if math.isnan(family['family_pearson']):
 family['family_pearson'] = 0
 matches = sorted(matches, key=lambda x: x['family_pearson'], reverse=True)
 matches = [family['superfeature'] for family in matches]
 matches = list(dict.fromkeys(matches))
 return gr.CheckboxGroup(choices=matches[:10]) # Limit to top 10 matches

 def visualize_feature_family(selected_families, current_subject):
 if not selected_families:
 return "Please select a feature family to visualize.", None, "", []

 selected_family = selected_families[0] # Take the first selected family
 family_analysis = subject_data[current_subject]['family_analysis']
 
 family_data = next((family for family in family_analysis if family['superfeature'] == selected_family), None)
 if not family_data:
 return "Invalid feature family selected.", None, "", []
 
 output = f"# {family_data['superfeature']}\n\n"
 
 # Create DataFrame
 df_data = [
 {
 "Feature": f"## {family_data['superfeature']}",
 "Parent Co-Occurrence": 1,
 "F1 Score": round(family_data['family_f1'], 2),
 "Pearson": round(family_data['family_pearson'], 4)
 },
 ]
 
 coocs = np.array(family_data['matrix'])[:, -1]
 # print(coocs)
 for name, cooc, f1, pearson in zip(family_data['feature_names'], coocs, family_data['feature_f1'], family_data['feature_pearson']):
 df_data.append({
 "Feature": name,
 "Parent Co-Occurrence": round(cooc, 2),
 "F1 Score": round(f1, 2),
 "Pearson": round(pearson, 4)
 })
 
 df = pd.DataFrame(df_data)
 
 # Add super reasoning below the dataframe
 output += "## Super Reasoning\n"
 output += f"{family_data['super_reasoning']}\n\n"
 
 graph = create_interactive_directed_graph(family_data)

 #return output, df, "", [], graph # Return empty string for search box and empty list for checkbox
 return df, "", [], graph

 family_search.change(search_feature_families, inputs=[family_search, subject], outputs=[family_matches])
 visualize_family_button.click(
 visualize_feature_family,
 inputs=[family_matches, subject],
 #outputs=[family_info, family_dataframe, family_search, family_matches, graph_plot]
 outputs=[family_dataframe, family_search, family_matches, graph_plot]
 )


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