Update app.py

app.py

@@ -125,12 +125,16 @@ examples = [
 
 topk = 5
 
-def get_features(text):
-
+# 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)
+    return sae_act
+
+def get_features(text):
+    sae_act = get_activations(text)
     sae_act_aggregated = ((sae_act[:,:,:] > 0).sum(1) > 0).cpu().numpy()
 
     X = pd.DataFrame(sae_act_aggregated)