Update app.py

app.py

@@ -69,6 +69,21 @@ def main(prob1):
     text = str(prob1).lower()
     obj = adr_predict(text)
     return obj[0], obj[1], obj[2]
+    
+# Define HTML for the legend
+legend_html = """
+<div style='margin-top: 20px;'>
+    <h3>NER Legend</h3>
+    <ul style='list-style-type:none;'>
+"""
+
+for entity, color in entity_colors.items():
+    legend_html += f"<li><span style='color:{color}; font-weight:bold;'>▉ </span>{entity}</li>"
+
+legend_html += "</ul></div>"
+
+# Create a Gradio HTML component to display the legend
+ner_legend = gr.HTML(value=legend_html)
 
 title = "Welcome to **ADR Tracker**"
 description1 = "This app takes text (up to a few sentences) and predicts to what extent the text describes severe (or non-severe) adverse reaction to medications. Please do NOT use for medical diagnosis."
@@ -100,7 +115,10 @@ with gr.Blocks(css=css) as demo:
         [label, local_plot, htext],
         api_name="adr"
     )
-    
+
+    # Display the NER Legend below the buttons
+    ner_legend  # Assuming you've defined this component above as shown
+
     with gr.Row():
         gr.Markdown("### Click on any of the examples below to see how it works:")
         gr.Examples([["A 35 year-old male had severe headache after taking Aspirin. The lab results were normal."],
@@ -108,3 +126,4 @@ with gr.Blocks(css=css) as demo:
                     [prob1], [label, local_plot, htext], main, cache_examples=True)
 
 demo.launch()
+