Update app.py

app.py

@@ -2,8 +2,6 @@ import json
 import streamlit as st
 from google.oauth2 import service_account
 from google.cloud import language_v1
-import urllib.parse
-import urllib.request
 import pandas as pd
 
 # Function to query Google's Knowledge Graph API
@@ -14,14 +12,9 @@ def query_knowledge_graph(entity_id):
     except Exception as e:
         st.write(f"An error occurred: {e}")
 
-# Function to count entities with 'mid' that contains '/g/' or '/m/' in their metadata
-def count_entities(entities):
-    count = 0
-    for entity in entities:
-        metadata = entity.metadata
-        if 'mid' in metadata and ('/g/' in metadata['mid'] or '/m/' in metadata['mid']):
-            count += 1
-    return count
+# Function to filter entities with "mid" in their metadata
+def filter_entities_with_mid(entities):
+    return [entity for entity in entities if 'mid' in entity.metadata]
 
 # Function to serialize entity metadata
 def serialize_entity_metadata(metadata):
@@ -80,22 +73,14 @@ def sample_analyze_entities(text_content):
 
     response = client.analyze_entities(request={"document": document, "encoding_type": encoding_type})
     
-    # Count the entities with 'mid' and either '/g/' or '/m/' in their metadata
-    entity_count = count_entities(response.entities)
-
-    if entity_count == 0:
-        st.markdown(f"# We found {len(response.entities)} entities - but found no Google Entities")
-        st.write("---")
-    elif entity_count == 1:
-        st.markdown(f"# We found {len(response.entities)} entities - and found 1 Google Entity")
-        st.write("---")
-    else:
-        st.markdown(f"# We found {len(response.entities)} entities - and found {entity_count} Google Entities")
-        st.write("---")
-
-
-    for i, entity in enumerate(response.entities):
-        st.write(f"Entity {i+1} of {len(response.entities)}")
+    # Filter entities that have a "mid" in their metadata
+    entities_with_mid = filter_entities_with_mid(response.entities)
+    
+    st.markdown(f"# We found {len(entities_with_mid)} entities with 'mid' in their metadata")
+    st.write("---")
+    
+    for i, entity in enumerate(entities_with_mid):
+        st.write(f"Entity {i+1} of {len(entities_with_mid)}")
         st.write(f"Name: {entity.name}")
         st.write(f"Type: {language_v1.Entity.Type(entity.type_).name}")
         st.write(f"Salience Score: {entity.salience}")
@@ -104,7 +89,7 @@ def sample_analyze_entities(text_content):
             st.write("Metadata:")
             st.write(entity.metadata)
             
-            if 'mid' in entity.metadata and ('/g/' in entity.metadata['mid'] or '/m/' in entity.metadata['mid']):
+            if 'mid' in entity.metadata:
                 entity_id = entity.metadata['mid']
                 query_knowledge_graph(entity_id)
 
@@ -118,7 +103,7 @@ def sample_analyze_entities(text_content):
         st.write("---")
 
     # Add the export functionality
-    export_entities(response.entities)
+    export_entities(entities_with_mid)
 
 # User input for text analysis
 user_input = st.text_area("Enter text to analyze")