Update app.py

app.py

@@ -31,6 +31,16 @@ ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-n
 ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
 #
 
+entity_colors = {
+    'Severity': '#E63946',  # a vivid red
+    'Sign_symptom': '#2A9D8F',  # a deep teal
+    'Medication': '#457B9D',  # a dusky blue
+    'Age': '#F4A261',  # a sandy orange
+    'Sex': '#F4A261',  # same sandy orange for consistency with 'Age'
+    'Diagnostic_procedure': '#9C6644',  # a brown
+    'Biological_structure': '#BDB2FF',  # a light pastel purple
+}
+
 def adr_predict(x):
     encoded_input = tokenizer(x, return_tensors='pt')
     output = model(**encoded_input)
@@ -41,15 +51,6 @@ def adr_predict(x):
     local_plot = shap.plots.text(shap_values[0], display=False)
 
     res = ner_pipe(x)
-    entity_colors = {
-        'Severity': '#E63946',  # a vivid red
-        'Sign_symptom': '#2A9D8F',  # a deep teal
-        'Medication': '#457B9D',  # a dusky blue
-        'Age': '#F4A261',  # a sandy orange
-        'Sex': '#F4A261',  # same sandy orange for consistency with 'Age'
-        'Diagnostic_procedure': '#9C6644',  # a brown
-        'Biological_structure': '#BDB2FF',  # a light pastel purple
-        }
 
     htext = ""
     prev_end = 0
@@ -65,6 +66,7 @@ def adr_predict(x):
    
     return {"Severe Reaction": float(scores.numpy()[1]), "Non-severe Reaction": float(scores.numpy()[0])}, local_plot, htext
 
+
 def main(prob1):
     text = str(prob1).lower()
     obj = adr_predict(text)