Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from annotated_text import annotated_text
 from refined.inference.processor import Refined
+import requests
 
 # Sidebar
 st.sidebar.image("logo-wordlift.png")
@@ -29,6 +30,14 @@ def get_wikidata_id(entity_string):
     entity_link = "https://www.wikidata.org/wiki/" + entity_id
     return {"id": entity_id, "link": entity_link}
 
+def get_entity_data(entity_link):
+    try:
+        response = requests.get(f'https://api.wordlift.io/id/{entity_link}')
+        return response.json()
+    except Exception as e:
+        print(f"Exception when fetching data for entity: {entity_link}. Exception: {e}")
+        return None
+    
 # Create the form
 with st.form(key='my_form'):
     text_input = st.text_area(label='Enter a sentence')
@@ -39,17 +48,16 @@ if text_input:
     entities = refined_model.process_text(text_input)
 
     entities_map = {}
-    entities_link_descriptions = {}
+    entities_data = {}
     for entity in entities:
         single_entity_list = str(entity).strip('][').replace("\'", "").split(', ')
         if len(single_entity_list) >= 2 and "wikidata" in single_entity_list[1]: 
             entities_map[single_entity_list[0].strip()] = get_wikidata_id(single_entity_list[1])
-            entities_link_descriptions[single_entity_list[0].strip()] = single_entity_list[2].strip().replace("(", "").replace(")", "")
+            entity_data = get_entity_data(entities_map[single_entity_list[0].strip()]["link"])
+            if entity_data is not None:
+                entities_data[single_entity_list[0].strip()] = entity_data
 
-    combined_entity_info_dictionary = dict([(k, [entities_map[k], entities_link_descriptions[k]]) for k in entities_map])
-
-    def get_entity_description(entity_string, combined_entity_info_dictionary):
-        return combined_entity_info_dictionary[entity_string][1]
+    combined_entity_info_dictionary = dict([(k, [entities_map[k], entities_data[k] if k in entities_data else None]) for k in entities_map])
 
     if submit_button:
         # Prepare a list to hold the final output
@@ -57,13 +65,26 @@ if text_input:
 
         # Replace each entity in the text with its annotated version
         for entity_string, entity_info in entities_map.items():
-            description = get_entity_description(entity_string, combined_entity_info_dictionary)
-            entity_annotation = (entity_string, entity_info["id"], "#8ef")  # Use the entity ID in the annotation
+            entity_data = entities_data.get(entity_string, None)
+            entity_type = None
+            if entity_data is not None:
+                entity_type = entity_data.get("@type", None)
+
+            # Use different colors based on the entity's type
+            color = "#8ef"  # Default color
+            if entity_type == "Place":
+                color = "#f00"
+            elif entity_type == "Organization":
+                color = "#0f0"
+            elif entity_type == "Person":
+                color = "#00f"
+
+            entity_annotation = (entity_string, entity_info["id"], color)
             text_input = text_input.replace(entity_string, f'{{{str(entity_annotation)}}}', 1)
 
         # Split the modified text_input into a list
         text_list = text_input.split("{")
-
+        
         for item in text_list:
             if "}" in item:
                 item_list = item.split("}")