biobert_ncbi_disease_ner

Model Description

This model is a fine-tuned version of BioBERT on the NCBI disease dataset for named entity recognition (NER) of diseases. It can be used to extract disease mentions from unstructured text in the medical and biological domains.

Intended Use

This model is intended for use in extracting disease mentions from unstructured text in the medical and biological domains. It can be used to improve information retrieval and knowledge extraction in these fields.

Training Data

This model was trained on the NCBI disease dataset, which consists of 793 PubMed abstracts with 6892 disease mentions.

How to use

You can use this model with the Hugging Face Transformers library. Here’s an example of how to load the model and use it to extract disease mentions from text:

from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline

tokenizer = AutoTokenizer.from_pretrained("ugaray96/biobert_ncbi_disease_ner")
model = AutoModelForTokenClassification.from_pretrained(
    "ugaray96/biobert_ncbi_disease_ner"
)

ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)

text = "The patient was diagnosed with lung cancer and started chemotherapy. They also have a history of diabetes and heart disease."
result = ner_pipeline(text)

diseases = []
for entity in result:
    if entity["entity"] == "Disease":
        diseases.append(entity["word"])
    elif entity["entity"] == "Disease Continuation" and diseases:
        diseases[-1] += f" {entity['word']}"

print(f"Diseases: {', '.join(diseases)}")

This should output: Diseases: lung cancer, diabetes, heart disease