Fix error prob bug

gec_model.py

@@ -78,6 +78,8 @@ class GecBERTModel(torch.nn.Module):
                 List of punctuations.
         """
         super().__init__()
+        if isinstance(model_paths, str):
+            model_paths = [model_paths]
         self.model_weights = list(map(float, weights)) if weights else [1] * len(model_paths)
         self.device = (
             torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else torch.device(device)
@@ -106,8 +108,6 @@ class GecBERTModel(torch.nn.Module):
 
         self.indexers = []
         self.models = []
-        if isinstance(model_paths, str):
-            model_paths = [model_paths]
         for model_path in model_paths:
             model = Seq2LabelsModel.from_pretrained(model_path)
             config = model.config
@@ -337,7 +337,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()