Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline
2
3model_name = "disham993/electrical-ner-ModernBERT-base"
4tokenizer = AutoTokenizer.from_pretrained(model_name)
5model = AutoModelForTokenClassification.from_pretrained(model_name)
6
7nlp = pipeline("ner", model=model, tokenizer=tokenizer, aggregation_strategy="simple")
8
9text = "The Xilinx Vivado development suite was used to program the Artix-7 FPGA."
10
11ner_results = nlp(text)
12
13def clean_and_group_entities(ner_results, min_score=0.40):
14 """
15 Cleans and groups named entity recognition (NER) results based on a minimum score threshold.
16
17 Args:
18 ner_results (list of dict): A list of dictionaries containing NER results. Each dictionary should have the keys:
19 - "word" (str): The recognized word or token.
20 - "entity_group" (str): The entity group or label.
21 - "start" (int): The start position of the entity in the text.
22 - "end" (int): The end position of the entity in the text.
23 - "score" (float): The confidence score of the entity recognition.
24 min_score (float, optional): The minimum score threshold for considering an entity. Defaults to 0.40.
25
26 Returns:
27 list of dict: A list of grouped entities that meet the minimum score threshold. Each dictionary contains:
28 - "entity_group" (str): The entity group or label.
29 - "word" (str): The concatenated word or token.
30 - "start" (int): The start position of the entity in the text.
31 - "end" (int): The end position of the entity in the text.
32 - "score" (float): The minimum confidence score of the grouped entity.
33 """
34 grouped_entities = []
35 current_entity = None
36
37 for result in ner_results:
38 # Skip entities with score below threshold
39 if result["score"] < min_score:
40 if current_entity:
41 # Add current entity if it meets threshold
42 if current_entity["score"] >= min_score:
43 grouped_entities.append(current_entity)
44 current_entity = None
45 continue
46
47 word = result["word"].replace("##", "") # Remove subword token markers
48
49 if current_entity and result["entity_group"] == current_entity["entity_group"] and result["start"] == current_entity["end"]:
50 # Continue the current entity
51 current_entity["word"] += word
52 current_entity["end"] = result["end"]
53 current_entity["score"] = min(current_entity["score"], result["score"])
54
55 # If combined score drops below threshold, discard the entity
56 if current_entity["score"] < min_score:
57 current_entity = None
58 else:
59 # Finalize the current entity if it meets threshold
60 if current_entity and current_entity["score"] >= min_score:
61 grouped_entities.append(current_entity)
62
63 # Start a new entity
64 current_entity = {
65 "entity_group": result["entity_group"],
66 "word": word,
67 "start": result["start"],
68 "end": result["end"],
69 "score": result["score"]
70 }
71
72 # Add the last entity if it meets threshold
73 if current_entity and current_entity["score"] >= min_score:
74 grouped_entities.append(current_entity)
75
76 return grouped_entities
77
78cleaned_results = clean_and_group_entities(ner_results)clean_and_group_entities function.@misc{modernbert,
title={Smarter, Better, Faster, Longer: A Modern Bidirectional Encoder for Fast, Memory Efficient, and Long Context Finetuning and Inference},
author={Benjamin Warner and Antoine Chaffin and Benjamin Clavié and Orion Weller and Oskar Hallström and Said Taghadouini and Alexis Gallagher and Raja Biswas and Faisal Ladhak and Tom Aarsen and Nathan Cooper and Griffin Adams and Jeremy Howard and Iacopo Poli},
year={2024},
eprint={2412.13663},
archivePrefix={arXiv},
primaryClass={cs.CL},
url={https://arxiv.org/abs/2412.13663},
}