Update handler.py

handler.py

@@ -22,14 +22,19 @@ class EndpointHandler():
         text = data.pop("text", data)
         tokens = self.tokenizer(text, return_tensors='pt', padding=True)
         output = self.model(**tokens)
-        vec = torch.max(
-            torch.log(
-                1 + torch.relu(output.logits)
-            ) * tokens.attention_mask.unsqueeze(-1),
-            dim=1)[0].squeeze()
-        cols = vec.nonzero().squeeze().cpu().tolist()
-        # extract the non-zero values
-        weights = vec[cols].cpu().tolist()
-        # use to create a dictionary of token ID to weight
-        sparse_dict = dict(zip(map(str, cols), weights))
-        return sparse_dict
+        results = []
+        for idx, x in enumerate(outputs.logits):
+            mask = tokens.attention_mask[idx]
+            mask = mask[None,:]
+            vec = torch.max(
+                torch.log(
+                    1 + torch.relu(x)
+                ) * mask.unsqueeze(-1),
+                dim=1)[0].squeeze()
+            cols = vec.nonzero().squeeze().cpu().tolist()
+            # extract the non-zero values
+            weights = vec[cols].cpu().tolist()
+            # use to create a dictionary of token ID to weight
+            sparse_dict = dict(zip(map(str, cols), weights))
+            results.append(sparse_dict)
+        return results