Update app.py

app.py

@@ -28,18 +28,25 @@ def count_entities(entities):
 def serialize_entity_metadata(metadata):
     return {k: str(v) for k, v in metadata.items()}
 
-# Function to export entities as a JSON or CSV file
+# Function to export entities as a JSON or CSV file, only exporting entities with 'mid' in their metadata
 def export_entities(entities):
     entity_list = []
     for entity in entities:
-        entity_info = {
-            "Name": entity.name,
-            "Type": language_v1.Entity.Type(entity.type_).name,
-            "Salience Score": entity.salience,
-            "Metadata": serialize_entity_metadata(entity.metadata),
-            "Mentions": [mention.text.content for mention in entity.mentions]
-        }
-        entity_list.append(entity_info)
+        # Check if entity has 'mid' in its metadata and if it contains '/g/' or '/m/' in 'mid'
+        if 'mid' in entity.metadata and ('/g/' in entity.metadata['mid'] or '/m/' in entity.metadata['mid']):
+            entity_info = {
+                "Name": entity.name,
+                "Type": language_v1.Entity.Type(entity.type_).name,
+                "Salience Score": entity.salience,
+                "Metadata": serialize_entity_metadata(entity.metadata),
+                "Mentions": [mention.text.content for mention in entity.mentions]
+            }
+            entity_list.append(entity_info)
+    
+    # If there are no entities to export, notify the user
+    if not entity_list:
+        st.write("No entities with a valid 'mid' found to export.")
+        return
     
     # Convert to DataFrame for easier export as CSV
     df = pd.DataFrame(entity_list)
@@ -64,8 +71,6 @@ st.sidebar.markdown("""
 5. **Export Entities**: Export the entities as JSON or CSV.
 """)
 
-
-
 # Header and intro
 st.title("Google Cloud NLP Entity Analyzer")
 st.write("This tool analyzes text to identify entities such as people, locations, organizations, and events.")