Update modeling_deberta.py

modeling_deberta.py

@@ -1158,6 +1158,55 @@ class DebertaV2ForMaskedLM(DebertaV2PreTrainedModel):
             attentions=outputs.attentions,
         )
 
+    @torch.no_grad()
+    def score(self, sequence: str, scored_length: int, tokenizer, device, batch_size):
+        mask_index = tokenizer.mask_token_id
+        cls_index = torch.tensor([tokenizer.cls_token_id])
+        sep_index = torch.tensor([tokenizer.sep_token_id])
+    
+        encoding = tokenizer(sequence, add_special_tokens=False, return_tensors="pt")
+        num_words = max(i for i in encoding.word_ids() if i is not None) + 1
+        scored_mask = [i and i >= num_words - scored_length for i in encoding.word_ids()]
+        num_to_score = sum(scored_mask)
+    
+        tokens = encoding.input_ids.squeeze(0)
+    
+        tokens = torch.cat([cls_index, tokens, sep_index]).to(device)
+        tokens = tokens.repeat(num_to_score, 1)
+        mask = torch.eye(tokens.size(1), device=device).bool()[-(num_to_score+1):-1, :]
+        input_ids = tokens.masked_fill(mask, value=mask_index)
+        if num_to_score > 1:
+            mask = torch.eye(tokens.size(1), device=device).bool()[-(num_to_score+1)+1:-1, :]
+            input_ids[:-1, :] = input_ids[:-1, :].masked_fill(mask, value=mask_index)
+        if num_to_score > 2:
+            mask = torch.eye(tokens.size(1), device=device).bool()[-(num_to_score+1)+2:-1, :]
+            input_ids[:-2, :] = input_ids[:-2, :].masked_fill(mask, value=mask_index)
+    
+        indices = torch.arange(input_ids.size(1) - num_to_score - 1, input_ids.size(1) - 1, device=device)
+        total_score = []
+    
+        for b in range((input_ids.size(0) - 1) // batch_size + 1):
+            logits = self(
+                input_ids[b * batch_size : (b+1) * batch_size, :].contiguous(),
+            ).logits
+    
+            logits = torch.gather(
+                logits,
+                dim=1,
+                index=indices[b * batch_size : (b+1) * batch_size].reshape(-1, 1, 1).expand(-1, -1, logits.size(-1))
+            ).squeeze(1)
+            log_p = F.log_softmax(logits, dim=-1)
+    
+            log_p = log_p.gather(
+                index=tokens[0, -(num_to_score+1):-1][b * batch_size : (b+1) * batch_size].unsqueeze(-1),
+                dim=-1
+            ).squeeze(-1)
+            total_score.append(log_p)
+    
+        total_score = torch.cat(total_score)
+        return total_score.sum().item()
+
+
 @add_start_docstrings("""DeBERTa Model with a `language modeling` head on top.""", DEBERTA_START_DOCSTRING)
 class DebertaV2ForCausalLM(DebertaV2ForMaskedLM):
     _tied_weights_keys = ["cls.predictions.decoder.weight", "cls.predictions.decoder.bias"]