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