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 = ["text"]

    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)

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

        # Extract entity labels and return as a list
        entity_labels = [entity["label"] for entity in entities]

        return entity_labels