🚀 DistilBert Drug NER by BlaikHole

📌 Overview

This repository provides a fine-tuned model trained on private drug NER data using quick still efficient DistilBert architecture. It can be used for Drug NER with 7 classes.

🎨 Model Outputs & Labels

The model identifies the following labels:

Label Name	Description
🟥 LABEL_0 > Dosage	Drug dosage/quantity mentioned in text.
🟩 LABEL_1 > Drug name	Drug name including Opioids.
🟦 LABEL_2 > Event	Event like overdose, death.
🟨 LABEL_3 > Location	Location of possible event.
🟪 LABEL_4 > Other (Outside)	Outside of rest 7 entities.
🟧 LABEL_5 > ROA	Route of Administration for drug usage like injection or oral use.
⬛ LABEL_6 > Symptom	Medical condition.
🟫 LABEL_7 > Temporal	Time, date or frequency.

🚀 Quick Usage

You can easily load and use this model with transformers:

🔹 Named Entity Recognition (NER)

import torch
from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer

# Label Mapping
LABEL_MAP = {
    0: "DOSAGE",
    1: "DRUG_NAME",
    2: "EVENT",
    3: "LOCATION",
    4: "OTHER",
    5: "ROA",
    6: "SYMPTOM",
    7: "TEMPORAL",
}

# Model Name
MODEL_NAME = "blaikhole/distilbert-drug-ner"

# Load Model & Tokenizer
device = 0 if torch.cuda.is_available() else -1
model = AutoModelForTokenClassification.from_pretrained(MODEL_NAME).to("cuda" if device == 0 else "cpu")
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# Load NER Pipeline
ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer, device=device)

def process_text(text):
    entities = ner_pipeline(text)
    results = [(entity["word"], LABEL_MAP.get(int(entity["entity"].split("_")[-1]), "OTHER")) for entity in entities]
    return results

# Example Usage
if __name__ == "__main__":
    sample_text = "The patient was prescribed oxycodone and experienced dizziness."
    print(process_text(sample_text))

📦 Installation

To use this model, install the required dependencies:

pip install transformers torch

📜 License

MIT

blaikhole
/

distilbert-drug-ner