Update ner_tool.py

ner_tool.py

@@ -1,3 +1,5 @@
+# Updated NamedEntityRecognitionTool in ner_tool.py
+
 from transformers import pipeline
 from transformers import Tool
 
@@ -18,6 +20,6 @@ class NamedEntityRecognitionTool(Tool):
         print(f"Identified Entities: {entities}")
 
         # Extract entity labels and return as a list
-        entity_labels = [entity["label"] for entity in entities]
+        entity_labels = [entity.get("entity", "UNKNOWN") for entity in entities]
 
-        return entity_labels
+        return {"entities": entity_labels}  # Return a dictionary with the specified output component