added confidence

generic_nel.py

@@ -120,16 +120,28 @@ class NelPipeline(Pipeline):
         wikipedia_predictions = self.tokenizer.batch_decode(
             outputs.sequences, skip_special_tokens=True
         )
-        print(len(outputs.sequences))
-        print(outputs.sequences.shape)
-        print(wikipedia_predictions)
-        print(len(scores))
-        for score in scores:
-            print(score.shape)
-        # Process scores and normalize
-        scores_tensor = scores.clone().detach()
-        probabilities = torch.exp(scores_tensor)
-        percentages = (probabilities * 100.0).cpu().numpy().tolist()
+        # Initialize confidence list for each token
+        token_confidences = []
+
+        # Process the scores for each token
+        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)
+
+        # 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