Update app.py

app.py

@@ -42,26 +42,16 @@ def get_channels(model):
     return None
 
 
-def forward_fn(input_ids, attention_mask, model, target_idx=None):
+def forward_fn(input_ids, attention_mask, model):
     output = model(input_ids=input_ids, attention_mask=attention_mask)
 
     if hasattr(output, "logits"):
-        logits = output.logits
-        if target_idx is not None and len(logits.shape) > 1 and logits.shape[-1] > 1:
-            return logits[:, target_idx]
-        return logits.squeeze(-1)
-
-    if hasattr(model, "encoder"):
-        encoder_output = model.encoder(
-            input_ids=input_ids, attention_mask=attention_mask
-        )
-        if hasattr(encoder_output, "last_hidden_state"):
-            return encoder_output.last_hidden_state[:, 0, :].mean(dim=-1)
-        return encoder_output[0][:, 0, :].mean(dim=-1)
+        return output.logits
+
+    if isinstance(output, tuple):
+        return output[0]
 
-    if hasattr(output, "last_hidden_state"):
-        return output.last_hidden_state[:, 0, :].mean(dim=-1)
-    return output[0][:, 0, :].mean(dim=-1)
+    return output
 
 
 @torch.no_grad()
@@ -93,17 +83,21 @@ def compute_attributions(
 
     baseline_ids = torch.full_like(input_ids, pad_id)
     lig = LayerIntegratedGradients(
-        lambda ids, am: forward_fn(ids, am, model, target_idx),
+        lambda ids, am: forward_fn(ids, am, model),
         get_embedding_layer(model),
     )
 
-    attributions, delta = lig.attribute(
-        inputs=input_ids,
-        baselines=baseline_ids,
-        additional_forward_args=(attention_mask,),
-        return_convergence_delta=True,
-        n_steps=n_steps,
-    )
+    attr_kwargs = {
+        "inputs": input_ids,
+        "baselines": baseline_ids,
+        "additional_forward_args": (attention_mask,),
+        "return_convergence_delta": True,
+        "n_steps": n_steps,
+    }
+    if target_idx is not None:
+        attr_kwargs["target"] = target_idx
+
+    attributions, delta = lig.attribute(**attr_kwargs)
 
     token_scores = attributions.sum(dim=-1) * attention_mask
     return token_scores, delta