feat: remote logic

app.py

@@ -1,12 +1,14 @@
 import gradio as gr
-import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from huggingface_hub import hf_hub_download
-import numpy as np
 import torch
-import pickle
 import numpy as np
 import pandas as pd
+import pickle
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
+import requests
+import os
+import msgpack_numpy as m
+import plotly.graph_objs as go
 from sklearn.linear_model import LogisticRegression
 
 torch.set_grad_enabled(False) # avoid blowing up mem
@@ -22,6 +24,8 @@ params = {
     "filename" : "layer_31/width_16k/average_l0_76/params.npz"
 }
 
+C = 0.01
+
 model_name = params["model_name"]
 width = params["width"]
 layer = params["layer"]
@@ -29,23 +33,17 @@ l0 = params["l0"]
 sae_repo_id = params["sae_repo_id"]
 filename = params["filename"]
 
-C = 0.01
-
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map='auto',
-    torch_dtype=torch.bfloat16,
-)
-tokenizer =  AutoTokenizer.from_pretrained(model_name)
-
 path_to_params = hf_hub_download(
     repo_id=sae_repo_id,
     filename=filename,
     force_download=False,
+    token=os.environ['TOKEN'],
 )
 
+tokenizer =  AutoTokenizer.from_pretrained(model_name)
+
 params = np.load(path_to_params)
-pt_params = {k: torch.from_numpy(v).cuda() for k, v in params.items()}
+pt_params = {k: torch.from_numpy(v) for k, v in params.items()}
 
 clf_name = f"linear_classifier_C_{C}_ "+ model_name + "_" + filename.split(".npz")[0]
 clf_name = clf_name.replace(os.sep, "_")
@@ -53,49 +51,6 @@ clf_name = clf_name.replace(os.sep, "_")
 with open(f"{clf_name}.pkl", 'rb') as model_file:
     clf: LogisticRegression = pickle.load(model_file)
 
-import torch.nn as nn
-class JumpReLUSAE(nn.Module):
-  def __init__(self, d_model, d_sae):
-    # Note that we initialise these to zeros because we're loading in pre-trained weights.
-    # If you want to train your own SAEs then we recommend using blah
-    super().__init__()
-    self.W_enc = nn.Parameter(torch.zeros(d_model, d_sae))
-    self.W_dec = nn.Parameter(torch.zeros(d_sae, d_model))
-    self.threshold = nn.Parameter(torch.zeros(d_sae))
-    self.b_enc = nn.Parameter(torch.zeros(d_sae))
-    self.b_dec = nn.Parameter(torch.zeros(d_model))
-
-  def encode(self, input_acts):
-    pre_acts = input_acts @ self.W_enc + self.b_enc
-    mask = (pre_acts > self.threshold)
-    acts = mask * torch.nn.functional.relu(pre_acts)
-    return acts
-
-  def decode(self, acts):
-    return acts @ self.W_dec + self.b_dec
-
-  def forward(self, acts):
-    acts = self.encode(acts)
-    recon = self.decode(acts)
-    return recon
-
-sae = JumpReLUSAE(params['W_enc'].shape[0], params['W_enc'].shape[1])
-sae.load_state_dict(pt_params)
-sae.to(dtype=torch.bfloat16).cuda()
-
-@torch.no_grad()
-def gather_residual_activations(model, target_layer, inputs):
-  target_act = None
-  def gather_target_act_hook(mod, inputs, outputs):
-    nonlocal target_act # make sure we can modify the target_act from the outer scope
-    target_act = outputs[0]
-    return outputs
-  handle = model.model.layers[target_layer].register_forward_hook(gather_target_act_hook)
-  _ = model.forward(inputs)
-  handle.remove()
-  return target_act
-
-import requests
 
 def get_feature_descriptions(feature):
     layer_name = f"{layer}-gemmascope-res-{width}"
@@ -127,15 +82,14 @@ topk = 5
 
 # Function to wrap in a FastAPI in case of
 def get_activations(text):
-    inputs = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True).to("cuda")
-
-    target_act = gather_residual_activations(model, layer, inputs)
-    sae_act = sae.encode(target_act)
+    response = requests.post("http://34.71.249.22:3000/execute_req", json={"query": text})
+    pack = m.unpackb(response.content)
+    sae_act = torch.from_numpy(pack["sae_act"]).to(dtype=torch.bfloat16)
     return sae_act
 
 def get_features(text):
     sae_act = get_activations(text)
-    sae_act_aggregated = ((sae_act[:,:,:] > 0).sum(1) > 0).cpu().numpy()
+    sae_act_aggregated = ((sae_act[:,:,:] > 0).sum(1) > 0).numpy()
 
     X = pd.DataFrame(sae_act_aggregated)
 
@@ -159,8 +113,6 @@ def get_features(text):
         descriptions.append(description)
     contrib_df["description"] = descriptions
 
-    import plotly.graph_objs as go
-
     fig = go.Figure(go.Bar(
         x=contrib_df['contribution'],
         y=contrib_df['description'],
@@ -190,17 +142,14 @@ def get_features(text):
     return classes, fig, dropdown
 
 def get_highlighted_text(text, feature):
-
-    inputs = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True).to("cuda")
-
-    target_act = gather_residual_activations(model, layer, inputs)
-    sae_act = sae.encode(target_act)
+    inputs = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)
+    sae_act = get_activations(text)
 
     activated_tokens = sae_act[0:,:,feature]
     max_activation = activated_tokens.max().item()
     activated_tokens /= max_activation
 
-    activated_tokens = activated_tokens.cpu().detach().float().numpy()
+    activated_tokens = activated_tokens.float().numpy()
 
     output = []
 

requirements.txt

@@ -0,0 +1,6 @@
+transformers
+numpy
+torch
+pandas
+scikit-learn
+msgpack-numpy