Update app.py

app.py

@@ -4,8 +4,8 @@ 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
 # Function to query Google's Knowledge Graph API
 def query_knowledge_graph(entity_id):
     try:
@@ -14,7 +14,6 @@ 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
@@ -24,6 +23,30 @@ def count_entities(entities):
             count += 1
     return count
 
+# Function to export entities as a JSON or CSV file
+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": entity.metadata,
+            "Mentions": [mention.text.content for mention in entity.mentions]
+        }
+        entity_list.append(entity_info)
+    
+    # Convert to DataFrame for easier export as CSV
+    df = pd.DataFrame(entity_list)
+    
+    # Export as CSV
+    csv = df.to_csv(index=False)
+    st.download_button(label="Export Entities as CSV", data=csv, file_name="entities.csv", mime="text/csv")
+    
+    # Export as JSON
+    json_data = json.dumps(entity_list, indent=2)
+    st.download_button(label="Export Entities as JSON", data=json_data, file_name="entities.json", mime="application/json")
+
 # Sidebar content
 st.sidebar.title("About This Tool")
 st.sidebar.markdown("This tool leverages Google's NLP technology for entity analysis.")
@@ -33,12 +56,13 @@ st.sidebar.markdown("""
 2. **User Input**: Enter the text you want to analyze.
 3. **Analyze**: Click the 'Analyze' button.
 4. **View Results**: See the identified entities and their details.
+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")
-st.write("Entity salience scores are always relative to the analysed text.")
+st.write("This tool analyzes text to identify entities such as people, locations, organizations, and events.")
+st.write("Entity salience scores are always relative to the analyzed text.")
 
 def sample_analyze_entities(text_content):
     service_account_info = json.loads(st.secrets["google_nlp"])
@@ -86,14 +110,14 @@ def sample_analyze_entities(text_content):
             st.write(f"Mentions: {mention_count} mention{plural}")
             st.write("Raw Array:")
             st.write(entity.mentions)
-            # st.write(', '.join([mention.text.content for mention in entity.mentions]))
-
 
         st.write("---")
 
+    # Add the export functionality
+    export_entities(response.entities)
+
 # User input for text analysis
 user_input = st.text_area("Enter text to analyze")
-#user_input = st.text_area("Enter text to analyze", max_chars=5000)
 
 if st.button("Analyze"):
     if user_input: