Update app.py

app.py

@@ -20,13 +20,15 @@ refined_model = load_model(selected_model_name)
 
 # Helper functions
 def get_wikidata_id(entity_string):
-  entity_list = entity_string.split("=")
-  return  "https://www.wikidata.org/wiki/" + str(entity_list[1])
+    entity_list = entity_string.split("=")
+    entity_id = str(entity_list[1])
+    entity_link = "https://www.wikidata.org/wiki/" + entity_id
+    return {"id": entity_id, "link": entity_link}
 
 # Create the form
 with st.form(key='my_form'):
     text_input = st.text_input(label='Enter a sentence')
-    submit_button = st.form_submit_button(label='Submit')
+    submit_button = st.form_submit_button(label='Analyze')
 
 # Process the text and extract the entities
 if text_input:
@@ -37,23 +39,23 @@ if text_input:
     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[get_wikidata_id(single_entity_list[1]).strip()] = single_entity_list[0].strip()
-            entities_link_descriptions[get_wikidata_id(single_entity_list[1]).strip()] = single_entity_list[2].strip().replace("(", "").replace(")", "")
+            entities_map[single_entity_list[0].strip()] = get_wikidata_id(single_entity_list[1]).strip()
+            entities_link_descriptions[single_entity_list[0].strip()] = single_entity_list[2].strip().replace("(", "").replace(")", "")
 
     combined_entity_info_dictionary = dict([(k, [entities_map[k], entities_link_descriptions[k]]) for k in entities_map])
 
-    def get_entity_description(entity_link, combined_entity_info_dictionary):
-        return combined_entity_info_dictionary[entity_link][1]
+    def get_entity_description(entity_string, combined_entity_info_dictionary):
+        return combined_entity_info_dictionary[entity_string][1]
 
     if submit_button:
         # Prepare a list to hold the final output
         final_text = []
 
         # Replace each entity in the text with its annotated version
-        for wikidata_link, entity in entities_map.items():
-            description = get_entity_description(wikidata_link, combined_entity_info_dictionary)
-            entity_annotation = (entity, description, "#8ef")
-            text_input = text_input.replace(entity, f'{{{str(entity_annotation)}}}', 1)
+        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
+            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("{")
@@ -69,3 +71,6 @@ if text_input:
 
         # Pass the final_text to the annotated_text function
         annotated_text(*final_text)
+
+with st.expander("See annotations"):
+    st.write(final_text)