Update app.py

app.py

@@ -8,19 +8,20 @@ nlp = spacy.load("./models/en_core_web_sm")
 st.title("SpaGAN Demo")
 st.write("Enter a text, and the system will highlight the geo-entities within it.")
 
-# Define a color map for different entity types
+# Define a color map and descriptions for different entity types
 COLOR_MAP = {
-    'FAC': 'red',    # Facilities (e.g., buildings, airports)
-    'ORG': 'blue',   # Organizations (e.g., companies, institutions)
-    'LOC': 'purple', # Locations (e.g., mountain ranges, water bodies)
-    'GPE': 'green'   # Geopolitical Entities (e.g., countries, cities)
+    'FAC': ('red', 'Facilities (e.g., buildings, airports)'),
+    'ORG': ('blue', 'Organizations (e.g., companies, institutions)'),
+    'LOC': ('purple', 'Locations (e.g., mountain ranges, water bodies)'),
+    'GPE': ('green', 'Geopolitical Entities (e.g., countries, cities)')
 }
 
-# Display the color key
+# Display the color key with descriptions
 st.write("**Color Key:**")
-for label, color in COLOR_MAP.items():
-    st.markdown(f"- **{label}**: <span style='color:{color}'>{color}</span>", unsafe_allow_html=True)
+for label, (color, description) in COLOR_MAP.items():
+    st.markdown(f"- **{label}**: <span style='color:{color}'>{color}</span> - {description}", unsafe_allow_html=True)
 
+# Text input
 user_input = st.text_area("Input Text", height=200)
 
 # Process the text when the button is clicked
@@ -33,7 +34,7 @@ if st.button("Highlight Geo-Entities"):
         highlighted_text = user_input
         for ent in reversed(doc.ents):
             if ent.label_ in COLOR_MAP:
-                color = COLOR_MAP[ent.label_]
+                color = COLOR_MAP[ent.label_][0]
                 highlighted_text = (
                     highlighted_text[:ent.start_char] +
                     f"<span style='color:{color}; font-weight:bold'>{ent.text}</span>" +