ner_tool.py

# Updated NamedEntityRecognitionTool in ner_tool.py

from transformers import pipeline
from transformers import Tool

class NamedEntityRecognitionTool(Tool):
    name = "ner_tool"
    description = "Identifies and labels entities such as persons, organizations, and locations in a given text."
    inputs = ["text"]
    outputs = ["entities"]

    def __call__(self, text: str):
        # Initialize the named entity recognition pipeline
        ner_analyzer = pipeline("ner")

        # Perform named entity recognition on the input text
        entities = ner_analyzer(text)

        # Extract relevant information for each identified entity
        entity_info = [{"entity": entity.get("entity", "UNKNOWN"), "word": entity.get("word", ""), "start": entity.get("start", -1), "end": entity.get("end", -1)} for entity in entities]

        # Extract the actual text span for each identified location entity
        location_entities = [text[start:end] for entity in entity_info if entity["entity"] == "I-LOC" for start, end in [(entity["start"], entity["end"])]]

        # Print the identified entities
        print(f"Identified Location Entities: {location_entities}")

        return {"entities": location_entities}  # Return a dictionary with the specified output component