fix probability logic

blimp.py

@@ -296,21 +296,22 @@ def get_batch_probabilities(
         with torch.no_grad():
             outputs = model(**inputs)
 
-        labels = inputs.input_ids
+        logits = outputs.logits[..., :-1, :].contiguous()
+        labels = inputs.input_ids[..., 1:].contiguous()
 
-        # Compute log probabilities
-        log_probs = torch.nn.functional.log_softmax(outputs.logits, dim=-1)
+        # compute log probabilities
+        log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
 
-        # Get probability of each actual token
+        # get per-token probability
         token_log_probs = torch.gather(log_probs, 2, labels.unsqueeze(-1)).squeeze(-1)
 
         if batch_size > 1:
-            # Create attention mask for padding
+            # mask padding tokens
             mask = (labels != tokenizer.pad_token_id).float()
             token_log_probs *= mask
 
         # sum log probabilities
-        sequence_log_probs = (token_log_probs).sum(dim=1)
+        sequence_log_probs = token_log_probs.sum(dim=1)
 
         probs.extend(sequence_log_probs.cpu().tolist())