added confidence

generic_nel.py

@@ -124,22 +124,24 @@ class NelPipeline(Pipeline):
         all_probabilities = []
         import torch.nn.functional as F
 
-        # 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]
+        print(type(scores))
+
+        # # 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