File size: 1,699 Bytes
cf00609 c5922b9 d25649c c5922b9 507724b c5922b9 cf00609 fb510e6 d25649c e607aa8 d25649c cf00609 d25649c cf00609 c5922b9 cf00609 c5922b9 cf00609 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
# Updated NamedEntityRecognitionTool in ner_tool.py
from transformers import pipeline
from transformers import Tool
class NamedEntityRecognitionTool(Tool):
name = "ner_tool"
description = "Identifies and labels various entities 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)
# Prepare a list to store token-level entities
token_entities = []
for entity in entities:
label = entity.get("entity", "UNKNOWN")
word = entity.get("word", "")
start = entity.get("start", -1)
end = entity.get("end", -1)
# Extract the complete entity text
entity_text = text[start:end].strip()
# Check for multi-token entities
if "##" in word:
# For multi-token entities, add each sub-token with its label
sub_tokens = word.split("##")
for i, sub_token in enumerate(sub_tokens):
token_entities.append({"token": sub_token, "label": label, "entity_text": entity_text})
else:
# For single-token entities, add the token with its label
token_entities.append({"token": word, "label": label, "entity_text": entity_text})
# Print the identified token-level entities
print(f"Token-level Entities: {token_entities}")
return {"entities": token_entities} # Return a dictionary with the specified output component
|