Fix error_prob bug

gec_model.py

@@ -89,6 +89,7 @@ class GecBERTModel(torch.nn.Module):
         self.lowercase_tokens = lowercase_tokens
         self.min_error_probability = min_error_probability
         self.vocab = Vocabulary.from_files(vocab_path)
+        self.incorr_index = self.vocab.get_token_index("INCORRECT", "d_tags")
         self.log = log
         self.iterations = iterations
         self.confidence = confidence
@@ -337,7 +338,10 @@ class GecBERTModel(torch.nn.Module):
         for output, weight in zip(data, self.model_weights):
             class_probabilities_labels = torch.softmax(output['logits'], dim=-1)
             all_class_probs += weight * class_probabilities_labels / sum(self.model_weights)
-            error_probs += weight * output['max_error_probability'] / sum(self.model_weights)
+            class_probabilities_d = torch.softmax(output['detect_logits'], dim=-1)
+            error_probs_d = class_probabilities_d[:, :, self.incorr_index]
+            incorr_prob = torch.max(error_probs_d, dim=-1)[0]
+            error_probs += weight * incorr_prob / sum(self.model_weights)
 
         max_vals = torch.max(all_class_probs, dim=-1)
         probs = max_vals[0].tolist()