Update app.py

app.py

@@ -3,15 +3,93 @@ from annotated_text import annotated_text
 import transformers
 
 ENTITY_TO_COLOR = {
-    'PER': '#8ef',
-    'LOC': '#faa',
-    'ORG': '#afa',
-    'MISC': '#fea',
+    'B-Activity': '#8ef',
+    'B-Administration': '#faa',
+    'B-Age': '#afa',
+    'B-Area': '#fea',
+    'B-Biological_attribute': '#8ef',
+    'B-Biological_structure': '#faa',
+    'B-Clinical_event': '#afa',
+    'B-Color': '#fea',
+    'B-Coreference': '#8ef',
+    'B-Date': '#faa',
+    'B-Detailed_description': '#afa',
+    'B-Diagnostic_procedure': '#fea',
+    'B-Disease_disorder': '#8ef',
+    'B-Distance': '#faa',
+    'B-Dosage': '#afa',
+    'B-Duration': '#fea',
+    'B-Family_history': '#8ef',
+    'B-Frequency': '#faa',
+    'B-Height': '#afa',
+    'B-History': '#fea',
+    'B-Lab_value': '#8ef',
+    'B-Mass': '#faa',
+    'B-Medication': '#afa',
+    'B-Nonbiological_location': '#fea',
+    'B-Occupation': '#8ef',
+    'B-Other_entity': '#faa',
+    'B-Other_event': '#afa',
+    'B-Outcome': '#fea',
+    'B-Personal_background': '#8ef',
+    'B-Qualitative_concept': '#faa',
+    'B-Quantitative_concept': '#afa',
+    'B-Severity': '#fea',
+    'B-Sex': '#8ef',
+    'B-Shape': '#faa',
+    'B-Sign_symptom': '#afa',
+    'B-Subject': '#fea',
+    'B-Texture': '#8ef',
+    'B-Therapeutic_procedure': '#faa',
+    'B-Time': '#afa',
+    'B-Volume': '#fea',
+    'B-Weight': '#8ef',
+    'I-Activity': '#faa',
+    'I-Administration': '#afa',
+    'I-Age': '#fea',
+    'I-Area': '#8ef',
+    'I-Biological_attribute': '#faa',
+    'I-Biological_structure': '#afa',
+    'I-Clinical_event': '#fea',
+    'I-Color': '#8ef',
+    'I-Coreference': '#faa',
+    'I-Date': '#afa',
+    'I-Detailed_description': '#fea',
+    'I-Diagnostic_procedure': '#8ef',
+    'I-Disease_disorder': '#faa',
+    'I-Distance': '#afa',
+    'I-Dosage': '#fea',
+    'I-Duration': '#8ef',
+    'I-Family_history': '#faa',
+    'I-Frequency': '#afa',
+    'I-Height': '#fea',
+    'I-History': '#8ef',
+    'I-Lab_value': '#faa',
+    'I-Mass': '#afa',
+    'I-Medication': '#fea',
+    'I-Nonbiological_location': '#8ef',
+    'I-Occupation': '#faa',
+    'I-Other_entity': '#afa',
+    'I-Other_event': '#fea',
+    'I-Outcome': '#8ef',
+    'I-Personal_background': '#faa',
+    'I-Qualitative_concept': '#afa',
+    'I-Quantitative_concept': '#fea',
+    'I-Severity': '#8ef',
+    'I-Shape': '#faa',
+    'I-Sign_symptom': '#afa',
+    'I-Subject': '#fea',
+    'I-Texture': '#8ef',
+    'I-Therapeutic_procedure': '#faa',
+    'I-Time': '#afa',
+    'I-Volume': '#fea',
+    'I-Weight': '#8ef',
+    'O': '#000'
 }
 
 @st.cache_data
 def get_pipe():
-    model_name = "dslim/bert-base-NER"
+    model_name = "nassga/nassBioMed"
     model = transformers.AutoModelForTokenClassification.from_pretrained(model_name)
     tokenizer = transformers.AutoTokenizer.from_pretrained(model_name)
     pipe = transformers.pipeline("token-classification", model=model, tokenizer=tokenizer, aggregation_strategy="simple")