added confidence

generic_nel.py

@@ -125,34 +125,12 @@ class NelPipeline(Pipeline):
         transition_scores = self.model.compute_transition_scores(
             outputs.sequences, outputs.scores, normalize_logits=True
         )
-        print(len(transition_scores), "!!!!")
-        probabilities = transition_scores[0]
-        print(probabilities, "!!!!", probabilities.shape)
         log_prob_sum = sum(transition_scores[0])
 
         # Calculate the probability for the entire sequence by exponentiating the sum of log probabilities
-
         sequence_confidence = torch.exp(log_prob_sum)
-
         percentages = sequence_confidence.cpu().numpy() * 100.0
 
-        # # Process each score (logits for the generated tokens)
-        # for i, score in enumerate(scores):
-        #     # Apply softmax to convert logits into probabilities
-        #     probabilities = F.softmax(score, dim=-1)
-        #
-        #     # Get the probabilities for the top tokens
-        #     top_probabilities = (
-        #         probabilities.cpu().numpy()
-        #     )  # Move to CPU and convert to NumPy
-        #
-        #     # Store the probabilities
-        #     all_probabilities.append(top_probabilities)
-        #     print(f"Top probabilities: {top_probabilities}")
-        #
-        # # Convert probabilities into percentages if needed
-        # percentages = [(prob * 100.0).tolist() for prob in all_probabilities]
-
         # Return the predictions along with the extracted entity, lOffset, and rOffset
         return wikipedia_predictions, enclosed_entity, lOffset, rOffset, [percentages]
 
@@ -196,7 +174,7 @@ class NelPipeline(Pipeline):
                     "wkd_id": qid,
                     "url": url,
                     "type": "UNK",
-                    "confidence_nel": percentages[idx],
+                    "confidence_nel": round(percentages[idx], 2),
                     "lOffset": lOffset,
                     "rOffset": rOffset,
                 }