commit files to HF hub

config.json

@@ -5,6 +5,15 @@
     "DistilBertForSequenceClassification"
   ],
   "attention_dropout": 0.1,
+  "custom_pipelines": {
+    "integrated-gradients": {
+      "impl": "integrated_gradients.IntegratedGradients",
+      "pt": [
+        "AutoModelForSequenceClassification"
+      ],
+      "tf": []
+    }
+  },
   "dim": 768,
   "dropout": 0.1,
   "hidden_dim": 3072,

integrated_gradients.py

@@ -0,0 +1,81 @@
+import transformers
+from transformers import (
+    Pipeline,
+    AutoTokenizer,
+    AutoModelForSequenceClassification
+)
+import torch
+import copy
+import time
+
+
+class IntegratedGradients(Pipeline):
+    def _sanitize_parameters(self, **kwargs):
+        forward_params = {}
+        if "baseline" in kwargs:
+            forward_params["baseline"] = kwargs["baseline"]
+        if "num_steps" in kwargs:
+            forward_params["num_steps"] = kwargs["num_steps"]
+        return {}, forward_params, {}
+
+    def preprocess(self, text):
+        return self.tokenizer(text, return_tensors=self.framework)
+
+    def _forward(
+        self,
+        model_inputs,
+        baseline,
+        num_steps
+    ):
+        torch.set_grad_enabled(True)
+        input_embed = self.model.base_model.embeddings.word_embeddings(model_inputs['input_ids'])
+        copy_embed = copy.deepcopy(input_embed.data)
+
+        if baseline is None:
+            # create baseline
+            baseline = torch.zeros_like(copy_embed)
+
+        grads = []
+
+        for step in range(num_steps + 1):
+            print(f"step: {step}/{num_steps}")
+            input_embed.data = baseline + step/num_steps * (copy_embed - baseline)
+            torch.set_grad_enabled(True)
+        input_embed = self.model.base_model.embeddings.word_embeddings(model_inputs['input_ids'])
+        copy_embed = copy.deepcopy(input_embed.data)
+
+        if baseline is None:
+            # create baseline
+            baseline = torch.zeros_like(copy_embed)
+
+        grads = []
+
+        for step in range(num_steps + 1):
+            print(f"step: {step}/{num_steps}")
+            input_embed.data = baseline + step/num_steps * (copy_embed - baseline)
+            print(input_embed.data)
+            outputs = self.model(**model_inputs, output_hidden_states=True, output_attentions=True)
+            logits, hidden_states = outputs.logits, outputs.hidden_states
+
+            # calculate the derivates of the output embeddings
+            out_embed = hidden_states[0]
+            g = torch.autograd.grad(logits, out_embed, grad_outputs=torch.ones_like(logits))[0]
+            grads.append(g)
+
+        # stack grads along first dimension to create a new tensor
+        grads = torch.stack(grads)
+
+        # approx integral
+        grads = (grads[:-1] + grads[1:]) / 2
+        avg_grad = grads.mean(0)
+
+        integrated_grads = out_embed * avg_grad
+
+        return integrated_grads
+
+    def postprocess(self, integrated_grads):
+        scores = torch.sqrt((integrated_grads ** 2).sum(-1))
+        # normalize scores
+        max_s, min_s = scores.max(1, True).values, scores.min(1, True).values
+        normalized_scores = (scores - min_s) / (max_s - min_s)
+        return normalized_scores[0].tolist()

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:19131abd6ebc0ccab0dc738d977ab9de31695047ec578fa8439704ed11f81a6a
+oid sha256:02a8411042d3ebcd765ee5ec4879eab02276f5747465af5ff2257d2fc492e2d1
 size 267832560