Update app.py

app.py

@@ -29,16 +29,25 @@ ner_tokenizer = AutoTokenizer.from_pretrained("d4data/biomedical-ner-all")
 ner_model = AutoModelForTokenClassification.from_pretrained("d4data/biomedical-ner-all")
 
 ner_pipe = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, aggregation_strategy="simple") # pass device=0 if using gpu
-#
+
+# fixing the colors
+# Create a custom color map
+cmap = {'0': '#457B9D',  # Non-severe reactions in blue
+        '1': '#E63946'}  # Severe reactions in red
+
 
 def adr_predict(x):
     encoded_input = tokenizer(x, return_tensors='pt')
     output = model(**encoded_input)
     scores = output[0][0].detach()
     scores = torch.nn.functional.softmax(scores)
-   
+
+    # Generate SHAP values and use the custom color map
     shap_values = explainer([str(x).lower()])
-    local_plot = shap.plots.text(shap_values[0], display=False)
+    
+    # Ensure the color depends on the output class; customize as needed
+    base_colors = {label: cmap[str(label)] for label in range(len(scores))} 
+    shap.plots.text(shap_values[0], color=base_colors, display=False)
 
     res = ner_pipe(x)
     entity_colors = {