Views
No views yet
bert-base-cased model for Medical Entity Extraction using the 'tner/bc5cdr' dataset. The model is specifically designed to recognize entities related to Disease,Symptoms,Drug. The model has been optimized for efficient deployment while maintaining high accuracy, making it suitable for resource-constrained environments.tner/bc5cdrpip install transformers torch1from transformers import BertTokenizerFast, BertForTokenClassification
2import torch
3
4device = "cuda" if torch.cuda.is_available() else "cpu"
5
6model_name = "AventIQ-AI/bert-medical-entity-extraction"
7model = BertForTokenClassification.from_pretrained(model_name).to(device)
8tokenizer = BertTokenizerFast.from_pretrained(model_name)1from transformers import pipeline
2
3ner_pipeline = pipeline("ner", model=model_name, tokenizer=tokenizer)
4test_sentence = "An overdose of Ibuprofen can lead to severe gastric issues."
5ner_results = ner_pipeline(test_sentence)
6label_map = {
7 "LABEL_0": "O", # Outside (not an entity)
8 "LABEL_1": "Drug",
9 "LABEL_2": "Disease",
10 "LABEL_3": "Symptom",
11 "LABEL_4": "Treatment"
12}
13
14def merge_tokens(ner_results):
15 merged_entities = []
16 current_word = ""
17 current_label = ""
18 current_score = 0
19 count = 0
20
21 for entity in ner_results:
22 word = entity["word"]
23 label = entity["entity"] # Model's output (e.g., LABEL_1, LABEL_2)
24 score = entity["score"]
25
26 # Merge subwords
27 if word.startswith("##"):
28 current_word += word[2:] # Remove '##' and append
29 current_score += score
30 count += 1
31 else:
32 if current_word: # Store the previous merged word
33 mapped_label = label_map.get(current_label, "Unknown")
34 merged_entities.append((current_word, mapped_label, current_score / count))
35 current_word = word
36 current_label = label
37 current_score = score
38 count = 1
39
40 # Add the last word
41 if current_word:
42 mapped_label = label_map.get(current_label, "Unknown")
43 merged_entities.append((current_word, mapped_label, current_score / count))
44
45 return merged_entities
46
47print("\n🩺 Medical NER Predictions:")
48for word, label, score in merge_tokens(ner_results):
49 if label != "O": # Skip non-entities
50 print(f"🔹 Entity: {word} | Category: {label} | Score: {score:.4f}")| Entity Type | Precision | Recall | F1 Score | Number of Entities |
|---|---|---|---|---|
| Disease | 91.46% | 92.07% | 91.76% | 3,000 |
| Drug | 71.25% | 72.83% | 72.03% | 1,266 |
| Symptom | 89.83% | 93.02% | 91.40% | 3,524 |
| Treatment | 88.83% | 92.02% | 90.40% | 3,124 |
tner/bc5cdr dataset was used, containing texts and their ner tags..
├── model/ # Contains the quantized model files
├── tokenizer_config/ # Tokenizer configuration and vocabulary files
├── model.safetensors/ # Quantized Model
├── README.md # Model documentation