Update app.py

app.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 import matplotlib
 from IPython.display import display, HTML
 from transformers import AutoTokenizer
-from DecompX.src.globenc_utils import GlobencConfig
+from DecompX.src.decompx_utils import DecompXConfig
 from DecompX.src.modeling_bert import BertForSequenceClassification
 from DecompX.src.modeling_roberta import RobertaForSequenceClassification
 
@@ -108,7 +108,7 @@ def run_decompx(text, model):
     SENTENCES = [text, "nothing"]
     CONFIGS = {
         "DecompX":
-            GlobencConfig(
+            DecompXConfig(
                 include_biases=True,
                 bias_decomp_type="absdot",
                 include_LN1=True,
@@ -144,12 +144,12 @@ def run_decompx(text, model):
     # RUN DECOMPX
     with torch.no_grad():
         model.eval()
-        logits, hidden_states, globenc_last_layer_outputs, globenc_all_layers_outputs = model(
+        logits, hidden_states, decompx_last_layer_outputs, decompx_all_layers_outputs = model(
             **tokenized_sentence, 
             output_attentions=False, 
             return_dict=False, 
             output_hidden_states=True, 
-            globenc_config=CONFIGS["DecompX"]
+            decompx_config=CONFIGS["DecompX"]
         )
     decompx_outputs = {
         "tokens": [tokenizer.convert_ids_to_tokens(tokenized_sentence["input_ids"][i][:batch_lengths[i]]) for i in range(len(SENTENCES))],
@@ -157,13 +157,13 @@ def run_decompx(text, model):
         "cls": hidden_states[-1][:, 0, :].cpu().detach().numpy().tolist()# Last layer & only CLS -> (batch, emb_dim)
     }
 
-    ### globenc_last_layer_outputs.classifier ~ (8, 55, 2) ###
-    importance = np.array([g.squeeze().cpu().detach().numpy() for g in globenc_last_layer_outputs.classifier]).squeeze()  # (batch, seq_len, classes)
+    ### decompx_last_layer_outputs.classifier ~ (8, 55, 2) ###
+    importance = np.array([g.squeeze().cpu().detach().numpy() for g in decompx_last_layer_outputs.classifier]).squeeze()  # (batch, seq_len, classes)
     importance = [importance[j][:batch_lengths[j], :] for j in range(len(importance))]
     decompx_outputs["importance_last_layer_classifier"] = importance
 
-    ### globenc_all_layers_outputs.aggregated ~ (12, 8, 55, 55) ###
-    importance = np.array([g.squeeze().cpu().detach().numpy() for g in globenc_all_layers_outputs.aggregated])  # (layers, batch, seq_len, seq_len)
+    ### decompx_all_layers_outputs.aggregated ~ (12, 8, 55, 55) ###
+    importance = np.array([g.squeeze().cpu().detach().numpy() for g in decompx_all_layers_outputs.aggregated])  # (layers, batch, seq_len, seq_len)
     importance = np.einsum('lbij->blij', importance)  # (batch, layers, seq_len, seq_len)
     importance = [importance[j][:, :batch_lengths[j], :batch_lengths[j]] for j in range(len(importance))]
     decompx_outputs["importance_all_layers_aggregated"] = importance
@@ -196,4 +196,4 @@ demo = gr.Interface(
     description="This is a demo for the ACL 2023 paper [DecompX](https://github.com/mohsenfayyaz/DecompX/)"
 )
 
-demo.launch()   
+demo.launch()