updated appliation

app.py

@@ -1,14 +1,268 @@
+import gradio as gr
 import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from safetensors import safe_open
+import os
+import json
+import math
+import random
+import numpy as np
+import matplotlib.pyplot as plt
+from graphviz import Digraph
+from PIL import Image, ImageDraw, ImageFont
+from io import BytesIO
+from sklearn.decomposition import PCA
 import logging
+import time
+from dotenv import load_dotenv
+from huggingface_hub import hf_hub_download
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Check if CUDA is available
+logger.info(f"HF_TOKEN_GEMMA set: {'HF_TOKEN_GEMMA' in os.environ}")
+logger.info(f"HF_TOKEN_EMBEDDINGS set: {'HF_TOKEN_EMBEDDINGS' in os.environ}")
+
 if not torch.cuda.is_available():
     raise RuntimeError("GPU is required but not available. ZeroGPU may not be initialized properly.")
 else:
     logger.info(f"CUDA is available. Device: {torch.cuda.get_device_name(0)}")
-    logger.info(f"Current CUDA device: {torch.cuda.current_device()}")
-    logger.info(f"Memory allocated: {torch.cuda.memory_allocated()} bytes")
+
+# Load environment variables
+load_dotenv()
+
+class Config:
+    def __init__(self):
+        self.MODEL_NAME = "google/gemma-2b"
+        self.ACCESS_TOKEN = os.environ.get("HF_TOKEN_GEMMA")
+        self.EMBEDDINGS_TOKEN = os.environ.get("HF_TOKEN_EMBEDDINGS")
+        self.DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+        self.DTYPE = torch.float32
+        self.TOPK = 5
+        self.CUTOFF = 0.00001  # Cumulative probability cutoff for tree branches
+        self.OUTPUT_LENGTH = 20
+        self.SUB_TOKEN_ID = 23070  # Arbitrary token ID to overwrite with embedding
+        self.LOG_BASE = 10
+
+config = Config()
+
+def load_tokenizer_and_model():
+    try:
+        logger.info(f"Loading tokenizer and model with token: {config.ACCESS_TOKEN[:5]}...")
+        tokenizer = AutoTokenizer.from_pretrained(config.MODEL_NAME, use_auth_token=config.ACCESS_TOKEN)
+        model = AutoModelForCausalLM.from_pretrained(config.MODEL_NAME, device_map="auto", use_auth_token=config.ACCESS_TOKEN)
+        model.to(config.DEVICE)  # Ensure the model is on the correct device
+        logger.info("Model and tokenizer loaded successfully")
+        return model, tokenizer
+    except Exception as e:
+        logger.error(f"Error loading model or tokenizer: {str(e)}")
+        return None, None
+
+def get_embeddings(model):
+    return model.get_input_embeddings().weight.data.to(config.DEVICE)
+
+
+def update_token_embedding(model, token_id, new_embedding):
+    new_embedding = new_embedding.to(model.get_input_embeddings().weight.device)
+    model.get_input_embeddings().weight.data[token_id] = new_embedding
+
+def produce_next_token_ids(input_ids, model, topk, sub_token_id):
+    input_ids = input_ids.to(model.device)
+    with torch.no_grad():
+        outputs = model(input_ids)
+        logits = outputs.logits
+    last_logits = logits[:, -1, :]
+    last_logits[:, sub_token_id] = float('-inf')
+    softmax_probs = torch.softmax(last_logits, dim=-1)
+    top_k_probs, top_k_ids = torch.topk(softmax_probs, k=topk, dim=-1)
+    return top_k_ids[0], top_k_probs[0]
+
+def build_def_tree(input_ids, data, base_prompt, model, tokenizer, config, depth=0, max_depth=25, cumulative_prob=1.0):
+    if depth >= max_depth or cumulative_prob < config.CUTOFF:
+        return
+
+    current_prompt = tokenizer.decode(input_ids[0], skip_special_tokens=True)
+    top_k_ids, top_k_probs = produce_next_token_ids(input_ids, model, config.TOPK, config.SUB_TOKEN_ID)
+
+    for idx, token_id in enumerate(top_k_ids.tolist()):
+        if token_id == config.SUB_TOKEN_ID:
+            continue  # Skip the substitute token to avoid circular definitions
+
+        token_id_tensor = torch.tensor([token_id], dtype=torch.long).to(model.device)
+        new_input_ids = torch.cat([input_ids, token_id_tensor.view(1, 1)], dim=-1)
+
+        new_cumulative_prob = cumulative_prob * top_k_probs[idx].item()
+
+        if new_cumulative_prob < config.CUTOFF:
+            continue
+
+        token_str = tokenizer.decode([token_id], skip_special_tokens=True)
+
+        new_child = {
+            "token_id": token_id,
+            "token": token_str,
+            "cumulative_prob": new_cumulative_prob,
+            "children": []
+        }
+        data['children'].append(new_child)
+
+        build_def_tree(new_input_ids, new_child, base_prompt, model, tokenizer, config, depth=depth+1, max_depth=max_depth, cumulative_prob=new_cumulative_prob)
+
+def generate_definition_tree(base_prompt, embedding, model, tokenizer, config):
+    results_dict = {"token": "", "cumulative_prob": 1, "children": []}
+
+    # Reset the token embedding
+    token_embedding = torch.unsqueeze(embedding, dim=0).to(model.device)
+    update_token_embedding(model, config.SUB_TOKEN_ID, token_embedding)
+
+    # Clear the model's cache if it has one
+    if hasattr(model, 'reset_cache'):
+        model.reset_cache()
+
+    input_ids = tokenizer.encode(base_prompt, return_tensors="pt").to(model.device)
+    build_def_tree(input_ids, results_dict, base_prompt, model, tokenizer, config)
+
+    return results_dict
+
+def find_max_min_cumulative_weight(node, current_max=0, current_min=float('inf')):
+    current_max = max(current_max, node.get('cumulative_prob', 0))
+    if node.get('cumulative_prob', 1) > 0:
+        current_min = min(current_min, node.get('cumulative_prob', 1))
+    for child in node.get('children', []):
+        current_max, current_min = find_max_min_cumulative_weight(child, current_max, current_min)
+    return current_max, current_min
+
+def scale_edge_width(cumulative_weight, max_weight, min_weight, log_base, max_thickness=33, min_thickness=1):
+    cumulative_weight = max(cumulative_weight, min_weight)
+    log_weight = math.log(cumulative_weight, log_base) - math.log(min_weight, log_base)
+    log_max = math.log(max_weight, log_base) - math.log(min_weight, log_base)
+    amplified_weight = (log_weight / log_max) ** 2.5
+    scaled_weight = (amplified_weight * (max_thickness - min_thickness)) + min_thickness
+    return scaled_weight
+
+def add_nodes_edges(dot, node, config, max_weight, min_weight, parent=None, is_root=True, depth=0, trim_cutoff=0):
+    node_id = str(id(node))
+    token = node.get('token', '').strip()
+    cumulative_prob = node.get('cumulative_prob', 1)
+
+    if cumulative_prob < trim_cutoff and not is_root:
+        return
+
+    if is_root or token:
+        if parent and not is_root:
+            edge_weight = scale_edge_width(cumulative_prob, max_weight, min_weight, config.LOG_BASE)
+            dot.edge(parent, node_id, arrowhead='dot', arrowsize='1', color='darkblue', penwidth=str(edge_weight))
+
+        label = "*" if is_root else token
+        dot.node(node_id, label=label, shape='plaintext', fontsize="36", fontname='Helvetica')
+
+        for child in node.get('children', []):
+            add_nodes_edges(dot, child, config, max_weight, min_weight, parent=node_id, is_root=False, depth=depth+1, trim_cutoff=trim_cutoff)
+
+def create_tree_diagram(data, config, trim_cutoff=0):
+    dot = Digraph(comment='Definition Tree', format='png')
+    dot.attr(rankdir='LR', size='5040,5000', margin='0.06', nodesep='0.06', ranksep='1', dpi='120', bgcolor='white')
+
+    max_weight, min_weight = find_max_min_cumulative_weight(data)
+    add_nodes_edges(dot, data, config, max_weight, min_weight, trim_cutoff=trim_cutoff)
+
+    output = BytesIO()
+    dot.render(outfile=output, format='png')
+    output.seek(0)
+
+    # Add white background
+    with Image.open(output) as img:
+        bg = Image.new("RGB", (img.width, 5000), (255, 255, 255))
+        y_offset = (5000 - img.height) // 2
+        bg.paste(img, (0, y_offset))
+        final_output = BytesIO()
+        bg.save(final_output, 'PNG')
+        final_output.seek(0)
+
+    return final_output
+
+def get_neuronpedia_url(layer, feature):
+    return f"https://neuronpedia.org/gemma-2b/{layer}-res-jb/{feature}?embed=true&embedexplanation=true&embedplots=false&embedtest=false&height=300"
+
+@torch.no_grad()
+def generate_definition_tree_placeholder(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight):
+    return "Definition tree generation placeholder"
+
+def gradio_interface():
+    model, tokenizer = load_tokenizer_and_model()
+    if model is None or tokenizer is None:
+        return gr.Interface(lambda: "Failed to load model and tokenizer. Please check the logs for more details.", inputs=[], outputs="text")
+
+    embeddings = get_embeddings(model)
+
+    with gr.Blocks() as demo:
+        gr.Markdown("# Gemma-2B SAE Feature Explorer")
+
+        with gr.Row():
+            with gr.Column():
+                selected_sae = gr.Dropdown(choices=["Gemma-2B layer 0", "Gemma-2B layer 6", "Gemma-2B layer 10", "Gemma-2B layer 12"], label="Select SAE")
+                feature_number = gr.Number(label="Select feature number", minimum=0, maximum=16383, value=0)
+
+                mode = gr.Radio(
+                    choices=["cosine distance token lists", "definition tree generation"],
+                    label="Select mode",
+                    value="cosine distance token lists"
+                )
+
+                weight_type = gr.Radio(["encoder", "decoder"], label="Select weight type for feature vector construction", value="encoder")
+                use_token_centroid = gr.Checkbox(label="Use token centroid offset", value=True)
+                scaling_factor = gr.Slider(minimum=0.1, maximum=10.0, value=3.8, label="Scaling factor (3.8 is mean distance from token embeddings to token centroid)")
+
+                num_exp = gr.Slider(minimum=0.1, maximum=5.0, value=1.4, label="Numerator exponent m")
+                denom_exp = gr.Slider(minimum=0.1, maximum=5.0, value=1.0, label="Denominator exponent n")
+                use_pca = gr.Checkbox(label="Introduce first PCA component")
+                pca_weight = gr.Slider(minimum=0.0, maximum=1.0, value=0.5, label="PCA weight")
+
+            with gr.Column():
+                output = gr.Image(label="Tree Diagram Output")
+                neuronpedia_embed = gr.HTML(label="Neuronpedia Embed")
+                trim_slider = gr.Slider(minimum=0.00001, maximum=0.1, value=0.00001, label="Trim cutoff for cumulative probability")
+
+        def update_output(selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode, trim_cutoff):
+            neuronpedia_url = get_neuronpedia_url(selected_sae.split(" ")[-1], feature_number)
+            neuronpedia_embed.update(value=f'<iframe src="{neuronpedia_url}" width="100%" height="300" frameborder="0"></iframe>')
+
+            if mode == "cosine distance token lists":
+                # Keep the original functionality here
+                pass
+            elif mode == "definition tree generation":
+                embedding = embeddings[int(feature_number)].to(config.DEVICE)
+                if use_token_centroid:
+                    token_centroid = torch.mean(embeddings, dim=0).to(config.DEVICE)
+                    embedding = token_centroid + scaling_factor * (embedding - token_centroid) / torch.norm(embedding - token_centroid)
+                
+                base_prompt = f'A typical definition of "{tokenizer.decode([config.SUB_TOKEN_ID], skip_special_tokens=True)}" would be "'
+                results_dict = generate_definition_tree(base_prompt, embedding, model, tokenizer, config)
+                tree_diagram = create_tree_diagram(results_dict, config, trim_cutoff=trim_cutoff)
+                return tree_diagram
+
+        def update_ui(mode_selected):
+            show_cosine_controls = mode_selected == "cosine distance token lists"
+            return (
+                gr.update(visible=show_cosine_controls),
+                gr.update(visible=show_cosine_controls),
+                gr.update(visible=show_cosine_controls),
+                gr.update(visible=show_cosine_controls)
+            )
+
+        inputs = [selected_sae, feature_number, weight_type, use_token_centroid, scaling_factor, use_pca, pca_weight, num_exp, denom_exp, mode, trim_slider]
+        mode.change(
+            update_ui, inputs=[mode], 
+            outputs=[num_exp, denom_exp, output]
+        )
+        
+        gr.Button("Generate Output").click(update_output, inputs=inputs, outputs=[output])
+
+    return demo
+
+
+if __name__ == "__main__":
+    iface = gradio_interface()
+    iface.launch()