Update app.py

app.py

@@ -2,8 +2,6 @@ import time
 import streamlit as st
 import pandas as pd
 import io
-
-from streamlit_extras.stylable_container import stylable_container
 import plotly.express as px
 import zipfile
 import os
@@ -15,6 +13,7 @@ from gliner import GLiNER
 from PyPDF2 import PdfReader 
 import docx 
 from comet_ml import Experiment 
+from streamlit_extras.stylable_container import stylable_container
 
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 
@@ -34,32 +33,89 @@ if 'file_upload_attempts' not in st.session_state:
 if 'encrypted_extracted_text' not in st.session_state:
     st.session_state['encrypted_extracted_text'] = None
 
-
-
 max_attempts = 10
 
 
-GLINER_LABELS = ["Person", "Organization", "Phone number", "Address", "Passport number", 
-"Email", "Credit card number", "Social security number", "Health insurance ID number", 
-"Date of birth", "Mobile phone number", "Bank account number", "Medication", "CPF", 
-"Driver license number", "Tax identification number", "Medical condition", 
-"Identity card number", "National ID number", "IP address", "IBAN", 
-"Credit card expiration date", "Username", "Health insurance number", 
-"Registration number", "Student ID number", "Insurance number", "Flight number", 
-"Landline phone number", "Blood type", "CVV", "Reservation number", 
-"Digital signature", "Social media handle", "License plate number", 
-"CNPJ", "Postal code", "Passport_number", "Serial number", "Vehicle registration number", 
-"Credit card brand", "Fax number", "Visa number", "Insurance company", "Identity document number", 
-"Transaction number", "National health insurance number", "CVC", "Birth certificate number", 
-"Train ticket number", "Passport expiration date", "Social_security_number"]
-
-    
-
+GLINER_LABELS_CATEGORIZED = {
+    "Personal Identifiers": [
+        "Person",
+        "Date of birth",
+        "Blood type",
+        "Digital signature",
+        "Social media handle",
+        "Username",
+        "Birth certificate number",
+    ],
+    "Contact Details": [
+        "Address",
+        "Phone number",
+        "Mobile phone number",
+        "Landline phone number",
+        "Email",
+        "Fax number",
+        "Postal code",
+    ],
+    "Financial & Payment": [
+        "Credit card number",
+        "Credit card expiration date",
+        "CVV",
+        "CVC",
+        "Bank account number",
+        "IBAN",
+        "Transaction number",
+        "Credit card brand",
+    ],
+    "Government & Official IDs": [
+        "Passport number",
+       
+        "Social security number",
+        
+        "CPF",
+        "Driver license number",
+        "Tax identification number",
+        "Identity card number",
+        "National ID number",
+        "Identity document number",
+        "Visa number",
+        "License plate number",
+        "CNPJ",
+        "Registration number",
+        "Student ID number",
+        "Passport expiration date",
+    ],
+    "Medical & Health": [
+        "Medication",
+        "Medical condition",
+        "Health insurance ID number",
+        "Health insurance number",
+        "National health insurance number",
+    ],
+    "Travel & Transport": [
+        "Flight number",
+        "Reservation number",
+        "Train ticket number",
+        "Vehicle registration number",
+    ],
+    "General Business & Other": [
+        "Organization",
+        "Insurance company",
+        "IP address",
+        "Serial number",
+        "Insurance number",
+    ]
+}
+
+# Flatten the categorized labels into a single list for GLiNER model input
+GLINER_LABELS_FLAT = [label for category_labels in GLINER_LABELS_CATEGORIZED.values() for label in category_labels]
+
+# Create a mapping from each specific label to its category for DataFrame processing
+LABEL_TO_CATEGORY_MAP = {label: category for category, labels in GLINER_LABELS_CATEGORIZED.items() for label in labels}
 
 @st.cache_resource
 def load_ner_model():
     """
-    Loads the pre-trained GLiNER NER model (urchade/gliner_multi_pii-v1) and caches it.
+    Loads the pre-trained GLiNER NER model (urchade/gliner_multi_pii-v1) and
+    caches it.
     This model is suitable for a wide range of custom entity types.
     """
     try:
@@ -116,34 +172,32 @@ def decrypt_text(encrypted_bytes: bytes) -> str | None:
 st.subheader("Multilingual PDF & DOCX Entity Finder", divider="orange") # Updated title
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
 
-expander = st.expander("**Important notes on the Multilingual PDF & DOCX Entity Finder**") # Updated title
+
+
+
+expander = st.expander("**Important notes on the Multilingual PDF & DOCX Entity Finder**") 
 expander.write(f'''
-    **Named Entities:** This Multilingual PDF & DOCX Entity Finder predicts a wide range of custom labels, including: {", ".join([f'"{label}"' for label in GLINER_LABELS])}.
-    
-    Results are presented in an easy-to-read table, visualized in an interactive tree map,
-    pie chart, and bar chart, and are available for download along with a Glossary of tags.
+    **Named Entities:** This Multilingual PDF & DOCX Entity Finder predicts a wide range of custom labels, including: "Person", "Organization", "Phone number", "Address", "Passport number", "Email", "Credit card number", "Social security number", "Health insurance ID number", "Date of birth", "Mobile phone number", "Bank account number", "Medication", "CPF", "Driver license number", "Tax identification number", "Medical condition", "Identity card number", "National ID number", "IP address", "IBAN", "Credit card expiration date", "Username", "Health insurance number", "Registration number", "Student ID number", "Insurance number", "Flight number", "Landline phone number", "Blood type", "CVV", "Reservation number", "Digital signature", "Social media handle", "License plate number", "CNPJ", "Postal code", "Serial number", "Vehicle registration number", "Credit card brand", "Fax number", "Visa number", "Insurance company", "Identity document number", "Transaction number", "National health insurance number", "CVC", "Birth certificate number", "Train ticket number", "Passport expiration date"
     
+    Results are presented in an easy-to-read table, visualized in an interactive tree map, pie chart, and bar chart, and are available for download along with a Glossary of tags.
     **Supported languages:** English, French, German, Spanish, Portuguese, Italian
     
-    **How to Use:** Upload your PDF or DOCX file. Then, click the 'Results' button
-    to extract and tag entities in your text data.
+    **How to Use:** Upload your PDF or DOCX file. Then, click the 'Results' button to extract and tag entities in your text data.
     
     **Usage Limits:** You can request results up to 10 times.
     
-    **Language settings:** Please check and adjust the language settings in
-    your computer, so the French, German, Spanish, Portuguese and Italian
-    characters are handled properly in your downloaded file.
+    **Language settings:** Please check and adjust the language settings in your computer, so the French, German, Spanish, Portuguese and Italian characters are handled properly in your downloaded file.
     
-    **Customization:** To change the app's background color to white or
-    black, click the three-dot menu on the right-hand side of your app, go to
-    Settings and then Choose app theme, colors and fonts.
+    **Customization:** To change the app's background color to white or black, click the three-dot menu on the right-hand side of your app, go to Settings and then Choose app theme, colors and fonts.
     
-    **Technical issues:** If your connection times out, please refresh the
-    page or reopen the app's URL.
+    **Technical issues:** If your connection times out, please refresh the page or reopen the app's URL.
     
     For any errors or inquiries, please contact us at [email protected]
 ''')
 
+    
+  
+
 with st.sidebar:
     container = st.container(border=True)
     container.write("**Named Entity Recognition (NER)** is the task of "
@@ -222,8 +276,8 @@ if st.button("Results"):
         
         # Measure NER model processing time
         start_time_ner = time.time()
-        # Use GLiNER's predict_entities method with the defined labels
-        text_entities = model.predict_entities(text_for_ner, GLINER_LABELS)
+        # Use GLiNER's predict_entities method with the defined flat list of labels
+        text_entities = model.predict_entities(text_for_ner, GLINER_LABELS_FLAT)
         end_time_ner = time.time()
         ner_processing_time = end_time_ner - start_time_ner
 
@@ -235,18 +289,18 @@ if st.button("Results"):
         else:
             st.error("Unexpected GLiNER output structure. Please check the model's output format.")
             st.stop()
-
         
-
         # Replace empty strings with 'Unknown' and drop rows with NaN after cleaning
         df = df.replace('', 'Unknown').dropna()
 
         if df.empty:
             st.warning("No entities were extracted from the uploaded text.")
             st.stop()
-
-        
         
+        # --- Add 'category' column to the DataFrame based on the grouped labels ---
+        df['category'] = df['entity_group'].map(LABEL_TO_CATEGORY_MAP)
+        # Handle cases where an entity_group might not have a category (shouldn't happen if maps are complete)
+        df['category'] = df['category'].fillna('Uncategorized') 
 
         if comet_initialized:
             experiment = Experiment(
@@ -266,7 +320,7 @@ if st.button("Results"):
         st.dataframe(df_styled, use_container_width=True)
 
         with st.expander("See Glossary of tags"):
-            st.write('''
+            st.write("""
             '**word**': ['entity extracted from your text data']
             
             '**score**': ['accuracy score; how accurately a tag has been assigned to
@@ -277,40 +331,42 @@ if st.button("Results"):
             '**start**': ['index of the start of the corresponding entity']
             
             '**end**': ['index of the end of the corresponding entity']
-            ''')
+            '**category**': ['the broader category this entity belongs to']
+            """)
 
         
-        st.subheader("Grouped entities", divider = "orange")
-
-        entity_items = [(label, label.replace('_', ' ').title()) for label in GLINER_LABELS]
-        tabs_per_row = 5
-        for i in range(0, len(entity_items), tabs_per_row):
-            current_row_entities = entity_items[i : i + tabs_per_row]
-            tab_titles = [item[1] for item in current_row_entities]
-
-            tabs = st.tabs(tab_titles)
-            for j, (entity_group_key, tab_title) in enumerate(current_row_entities):
-                with tabs[j]:
-                    if entity_group_key in df["entity_group"].unique():
-                        df_filtered = df[df["entity_group"] == entity_group_key]
-                        st.dataframe(df_filtered, use_container_width=True)
-                    else:
-                        st.info(f"No '{tab_title}' entities found in the text.")
-                        # Display an empty DataFrame for consistency if no entities are found
-                        st.dataframe(pd.DataFrame({
-                            'entity_group': [entity_group_key],
-                            'score': [np.nan],
-                            'word': [np.nan],
-                            'start': [np.nan],
-                            'end': [np.nan]
-                        }), hide_index=True)
+        st.subheader("Grouped Entities by Category", divider = "orange")
+        
+        # Create tabs for each category
+        category_names = list(GLINER_LABELS_CATEGORIZED.keys())
+        category_tabs = st.tabs(category_names)
+
+        for i, category_name in enumerate(category_names):
+            with category_tabs[i]:
+                
+                
+                # Filter the main DataFrame for the current category
+                df_category_filtered = df[df['category'] == category_name]
+
+                if not df_category_filtered.empty:
+                    # Sort entities within the category by their specific type for better display
+                    for entity_type in GLINER_LABELS_CATEGORIZED[category_name]:
+                        df_entity_type_filtered = df_category_filtered[df_category_filtered['entity_group'] == entity_type]
+                        if not df_entity_type_filtered.empty:
+                            st.markdown(f"***{entity_type}***")
+                            st.dataframe(df_entity_type_filtered.drop(columns=['category']), use_container_width=True)
+                        else:
+                            st.info(f"No '{entity_type}' entities found for this category.")
+                else:
+                    st.info(f"No entities found for the '{category_name}' category.")
 
         st.divider()
 
         # --- Visualizations ---
         st.subheader("Tree map", divider="orange")
-        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'entity_group', 'word'], # Changed path for better visual grouping
-                                 values='score', color='entity_group')
+        # Update treemap path to include category
+        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'category', 'entity_group', 'word'],
+                                 values='score', color='category') # Color by category for better visual distinction
         fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25))
         st.plotly_chart(fig_treemap)
         if comet_initialized:
@@ -321,32 +377,43 @@ if st.button("Results"):
 
         col1, col2 = st.columns(2)
         with col1:
-            st.subheader("Pie Chart", divider="orange")
+            st.subheader("Pie Chart (by Entity Type)", divider="orange")
             fig_pie = px.pie(final_df_counts, values='count', names='entity_group',
-                             hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted labels')
+                             hover_data=['count'], labels={'count': 'count'}, title='Percentage of Predicted Labels (Entity Types)')
             fig_pie.update_traces(textposition='inside', textinfo='percent+label')
             st.plotly_chart(fig_pie)
             if comet_initialized:
                 experiment.log_figure(figure=fig_pie, figure_name="label_pie_chart")
 
         with col2:
-            st.subheader("Bar Chart", divider="orange")
+            st.subheader("Bar Chart (by Entity Type)", divider="orange")
             fig_bar = px.bar(final_df_counts, x="count", y="entity_group", color="entity_group", text_auto=True,
-                             title='Occurrences of predicted labels')
+                             title='Occurrences of Predicted Labels (Entity Types)', orientation='h')
+            fig_bar.update_layout(yaxis={'categoryorder':'total ascending'}) # Order bars
             st.plotly_chart(fig_bar)
             if comet_initialized:
                 experiment.log_figure(figure=fig_bar, figure_name="label_bar_chart")
 
+        # Add a chart for categories
+        st.subheader("Entity Counts by Category", divider="orange")
+        category_counts = df['category'].value_counts().reset_index().rename(columns={"index": "category", "count": "count"})
+        fig_cat_bar = px.bar(category_counts, x="count", y="category", color="category", text_auto=True,
+                             title='Occurrences of Entities by Category', orientation='h')
+        fig_cat_bar.update_layout(yaxis={'categoryorder':'total ascending'})
+        st.plotly_chart(fig_cat_bar)
+
+
         # --- Downloadable Content ---
         dfa = pd.DataFrame(
             data={
-                'Column Name': ['word', 'entity_group','score', 'start', 'end'],
+                'Column Name': ['word', 'entity_group', 'score', 'start', 'end', 'category'],
                 'Description': [
                     'entity extracted from your text data',
                     'label (tag) assigned to a given extracted entity',
                     'accuracy score; how accurately a tag has been assigned to a given entity',
                     'index of the start of the corresponding entity',
                     'index of the end of the corresponding entity',
+                    'the broader category this entity belongs to',
                 ]
             }
         )