Views
No views yet
1from transformers import AutoTokenizer, AutoModelForTokenClassification
2from transformers import pipeline
3
4tokenizer = AutoTokenizer.from_pretrained("ugaray96/biobert_ncbi_disease_ner")
5model = AutoModelForTokenClassification.from_pretrained(
6 "ugaray96/biobert_ncbi_disease_ner"
7)
8
9ner_pipeline = pipeline("ner", model=model, tokenizer=tokenizer)
10
11text = "The patient was diagnosed with lung cancer and started chemotherapy. They also have a history of diabetes and heart disease."
12result = ner_pipeline(text)
13
14diseases = []
15for entity in result:
16 if entity["entity"] == "Disease":
17 diseases.append(entity["word"])
18 elif entity["entity"] == "Disease Continuation" and diseases:
19 diseases[-1] += f" {entity['word']}"
20
21print(f"Diseases: {', '.join(diseases)}")Diseases: lung cancer, diabetes, heart disease