Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,6 +12,7 @@ from streamlit_extras.stylable_container import stylable_container
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
+
 st.markdown(
     """
     <style>
@@ -51,7 +52,9 @@ st.markdown(
     }
     </style>
     """,
-    unsafe_allow_html=True)
+    unsafe_allow_html=True
+)
+
 # --- Page Configuration and UI Elements ---
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 st.subheader("PiiGuard", divider="violet")
@@ -59,17 +62,18 @@ st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
 expander = st.expander("**Important notes**")
 expander.write("""**Named Entities:** This PiiGuard web app predicts fifty-one (51) labels: "person", "organization", "social_media_handle", "username", "insurance_company", "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number", "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc", "bank_account_number", "iban", "transaction_number", "cpf", "cnpj",  "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number", "identity_card_number", "national_id_number", "identity_document_number", "birth_certificate_number", "social_security_number", "health_insurance_id_number", "health_insurance_number", "national_health_insurance_number", "student_id_number", "registration_number", "insurance_number", "serial_number", "visa_number", "reservation_number", "train_ticket_number", "medication", "medical_condition", "blood_type",  "date_of_birth", "address", "ip_address", "postal_code",  "flight_number", "license_plate_number", "vehicle_registration_number", "digital_signature"
 
-Results are presented in easy-to-read tables, visualized in an interactive tree map, pie chart and bar chart, and are available for download along with a Glossary of tags.      
+Results are presented in easy-to-read tables, visualized in an interactive tree map, pie chart and bar chart, and are available for download along with a Glossary of tags.
 
-**How to Use:** Type or paste your text into the text area below, then press Ctrl + Enter. Click the 'Results' button to extract and tag entities in your text data.      
+**How to Use:** Type or paste your text into the text area below, then press Ctrl + Enter. Click the 'Results' button to extract and tag entities in your text data.
 
 **Usage Limits:** You can request results unlimited times for one (1) month.
 
-**Supported Languages:** English 
+**Supported Languages:** English
 
 **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 info@nlpblogs.com""")
+
 with st.sidebar:
     st.write("Use the following code to embed the PiiGuard web app on your website. Feel free to adjust the width and height values to fit your page.")
     code = '''
@@ -87,6 +91,7 @@ with st.sidebar:
     st.divider()
     st.subheader("🚀 Ready to build your own AI Web App?", divider="violet")
     st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", type="primary")
+
 # --- Comet ML Setup ---
 COMET_API_KEY = os.environ.get("COMET_API_KEY")
 COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
@@ -94,68 +99,20 @@ COMET_PROJECT_NAME = os.environ.get("COMET_PROJECT_NAME")
 comet_initialized = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME)
 if not comet_initialized:
     st.warning("Comet ML not initialized. Check environment variables.")
+
 # --- Label Definitions ---
-labels = [ "person",
-    "organization",
-    "social_media_handle",
-    "username",
-    "insurance_company", "phone_number",
-    "email",
-    "email_address",
-    "mobile_phone_number",
-    "landline_phone_number",
-    "fax_number", "credit_card_number",
-    "credit_card_expiration_date",
-    "credit_card_brand",
-    "cvv",
-    "cvc",
-    "bank_account_number",
-    "iban",
-    "transaction_number",
-    "cpf",
-    "cnpj",  "passport_number",
-    "passport_expiration_date",
-    "driver's_license_number",
-    "tax_identification_number",
-    "identity_card_number",
-    "national_id_number",
-    "identity_document_number",
-    "birth_certificate_number",
-    "social_security_number",
-    "health_insurance_id_number",
-    "health_insurance_number",
-    "national_health_insurance_number",
-    "student_id_number",
-    "registration_number",
-    "insurance_number",
-    "serial_number",
-    "visa_number",
-    "reservation_number",
-    "train_ticket_number", "medication",
-    "medical_condition",
-    "blood_type",
-    "date_of_birth", "address",
-    "ip_address",
-    "postal_code",  "flight_number",
-    "license_plate_number",
-    "vehicle_registration_number", "digital_signature"]
-# Corrected mapping dictionary
-category_mapping = {  "People_and_Groups": [
+labels = [
     "person",
     "organization",
     "social_media_handle",
     "username",
-    "insurance_company"
-  ],
-  "Contact_Information": [
+    "insurance_company",
     "phone_number",
     "email",
     "email_address",
     "mobile_phone_number",
     "landline_phone_number",
-    "fax_number"
-  ],
-  "Financial_and_Transactions": [
+    "fax_number",
     "credit_card_number",
     "credit_card_expiration_date",
     "credit_card_brand",
@@ -165,9 +122,7 @@ category_mapping = {  "People_and_Groups": [
     "iban",
     "transaction_number",
     "cpf",
-    "cnpj"
-  ],
-  "Identification_and_Documents": [
+    "cnpj",
     "passport_number",
     "passport_expiration_date",
     "driver's_license_number",
@@ -186,164 +141,276 @@ category_mapping = {  "People_and_Groups": [
     "serial_number",
     "visa_number",
     "reservation_number",
-    "train_ticket_number"
-  ],
-  "Health_and_Personal": [
+    "train_ticket_number",
     "medication",
     "medical_condition",
     "blood_type",
-    "date_of_birth"
-  ],
-  "Locations_and_Addresses": [
+    "date_of_birth",
     "address",
     "ip_address",
-    "postal_code"
-  ],
-  "Transportation_and_Logistics": [
+    "postal_code",
     "flight_number",
     "license_plate_number",
-    "vehicle_registration_number"
-  ],
-  "Digital_and_Security": [
+    "vehicle_registration_number",
     "digital_signature"
-  ]}
+]
+# Corrected mapping dictionary
+category_mapping = {
+    "People_and_Groups": [
+        "person",
+        "organization",
+        "social_media_handle",
+        "username",
+        "insurance_company"
+    ],
+    "Contact_Information": [
+        "phone_number",
+        "email",
+        "email_address",
+        "mobile_phone_number",
+        "landline_phone_number",
+        "fax_number"
+    ],
+    "Financial_and_Transactions": [
+        "credit_card_number",
+        "credit_card_expiration_date",
+        "credit_card_brand",
+        "cvv",
+        "cvc",
+        "bank_account_number",
+        "iban",
+        "transaction_number",
+        "cpf",
+        "cnpj"
+    ],
+    "Identification_and_Documents": [
+        "passport_number",
+        "passport_expiration_date",
+        "driver's_license_number",
+        "tax_identification_number",
+        "identity_card_number",
+        "national_id_number",
+        "identity_document_number",
+        "birth_certificate_number",
+        "social_security_number",
+        "health_insurance_id_number",
+        "health_insurance_number",
+        "national_health_insurance_number",
+        "student_id_number",
+        "registration_number",
+        "insurance_number",
+        "serial_number",
+        "visa_number",
+        "reservation_number",
+        "train_ticket_number"
+    ],
+    "Health_and_Personal": [
+        "medication",
+        "medical_condition",
+        "blood_type",
+        "date_of_birth"
+    ],
+    "Locations_and_Addresses": [
+        "address",
+        "ip_address",
+        "postal_code"
+    ],
+    "Transportation_and_Logistics": [
+        "flight_number",
+        "license_plate_number",
+        "vehicle_registration_number"
+    ],
+    "Digital_and_Security": [
+        "digital_signature"
+    ]
+}
+
 # --- Model Loading ---
-@st.cache_resourcedef load_ner_model():
+@st.cache_resource
+def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
-        return GLiNER.from_pretrained("knowledgator/gliner-multitask-v1.0", nested_ner=True, num_gen_sequences=2, gen_constraints= labels)
+        return GLiNER.from_pretrained("knowledgator/gliner-multitask-v1.0", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
     except Exception as e:
         st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
         st.stop()
 model = load_ner_model()
+
 # Flatten the mapping to a single dictionary
 reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}
+
+# --- Session State Initialization ---
+if 'show_results' not in st.session_state:
+    st.session_state.show_results = False
+if 'last_text' not in st.session_state:
+    st.session_state.last_text = ""
+if 'results_df' not in st.session_state:
+    st.session_state.results_df = pd.DataFrame()
+if 'elapsed_time' not in st.session_state:
+    st.session_state.elapsed_time = 0.0
+
 # --- Text Input and Clear Button ---
 word_limit = 200
 text = st.text_area(f"Type or paste your text below (max {word_limit} words), and then press Ctrl + Enter", height=250, key='my_text_area')
 word_count = len(text.split())
 st.markdown(f"**Word count:** {word_count}/{word_limit}")
+
 def clear_text():
-    """Clears the text area."""
+    """Clears the text area and hides results."""
     st.session_state['my_text_area'] = ""
+    st.session_state.show_results = False
+    st.session_state.last_text = ""
+    st.session_state.results_df = pd.DataFrame()
+    st.session_state.elapsed_time = 0.0
 st.button("Clear text", on_click=clear_text)
+
 # --- Results Section ---
 if st.button("Results"):
-    start_time = time.time()
     if not text.strip():
         st.warning("Please enter some text to extract entities.")
+        st.session_state.show_results = False
     elif word_count > word_limit:
         st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
+        st.session_state.show_results = False
     else:
-        with st.spinner("Extracting entities...", show_time=True):
-            entities = model.predict_entities(text, labels)
-            df = pd.DataFrame(entities)
-            if not df.empty:
-                df['category'] = df['label'].map(reverse_category_mapping)
-                if comet_initialized:
-                    experiment = Experiment(
-                        api_key=COMET_API_KEY,
-                        workspace=COMET_WORKSPACE,
-                        project_name=COMET_PROJECT_NAME,
-                    )
-                    experiment.log_parameter("input_text", text)
-                    experiment.log_table("predicted_entities", df)
-                st.subheader("Grouped Entities by Category", divider = "violet")
-                # Create tabs for each category
-                category_names = sorted(list(category_mapping.keys()))
-                category_tabs = st.tabs(category_names)
-                for i, category_name in enumerate(category_names):
-                    with category_tabs[i]:
-                        df_category_filtered = df[df['category'] == category_name]
-                        if not df_category_filtered.empty:
-                            st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
-                        else:
-                            st.info(f"No entities found for the '{category_name}' category.")
-                with st.expander("See Glossary of tags"):
-                    st.write('''
-                    - **text**: ['entity extracted from your text data']
-                    - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
-                    - **label**: ['label (tag) assigned to a given extracted entity']
-                    - **start**: ['index of the start of the corresponding entity']
-                    - **end**: ['index of the end of the corresponding entity']
-                    ''')
-                st.divider()
-                # Tree map
-                st.subheader("Tree map", divider = "violet")
-                fig_treemap = px.treemap(df, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
-                fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#E8F5E9', plot_bgcolor='#E8F5E9')
-                st.plotly_chart(fig_treemap)
-                # Pie and Bar charts
-                grouped_counts = df['category'].value_counts().reset_index()
-                grouped_counts.columns = ['category', 'count']
-                col1, col2 = st.columns(2)
-                with col1:
-                    st.subheader("Pie chart", divider = "violet")
-                    fig_pie = px.pie(grouped_counts, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
-                    fig_pie.update_traces(textposition='inside', textinfo='percent+label')
-                    fig_pie.update_layout(
-                        paper_bgcolor='#E8F5E9',
-                        plot_bgcolor='#E8F5E9'
-                    )
-                    st.plotly_chart(fig_pie)
-                with col2:
-                    st.subheader("Bar chart", divider = "violet")
-                    fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
-                    fig_bar.update_layout(  # Changed from fig_pie to fig_bar
-                        paper_bgcolor='#E8F5E9',
-                        plot_bgcolor='#E8F5E9'
-                    )
-                    st.plotly_chart(fig_bar)
-                # Most Frequent Entities
-                st.subheader("Most Frequent Entities", divider="gray")
-                word_counts = df['text'].value_counts().reset_index()
-                word_counts.columns = ['Entity', 'Count']
-                repeating_entities = word_counts[word_counts['Count'] > 1]
-                if not repeating_entities.empty:
-                    st.dataframe(repeating_entities, use_container_width=True)
-                    fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
-                    fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'},
-                        paper_bgcolor='#F5F5F5',
-                        plot_bgcolor='#F5F5F5')
-                    st.plotly_chart(fig_repeating_bar)
+        # Check if the text is different from the last time
+        if text != st.session_state.last_text:
+            st.session_state.show_results = True
+            st.session_state.last_text = text
+            start_time = time.time()
+            with st.spinner("Extracting entities...", show_time=True):
+                entities = model.predict_entities(text, labels)
+                df = pd.DataFrame(entities)
+                st.session_state.results_df = df
+                if not df.empty:
+                    df['category'] = df['label'].map(reverse_category_mapping)
+                    if comet_initialized:
+                        experiment = Experiment(
+                            api_key=COMET_API_KEY,
+                            workspace=COMET_WORKSPACE,
+                            project_name=COMET_PROJECT_NAME,
+                        )
+                        experiment.log_parameter("input_text", text)
+                        experiment.log_table("predicted_entities", df)
+                        experiment.end()
+            end_time = time.time()
+            st.session_state.elapsed_time = end_time - start_time
+        else:
+            # If the text is the same, just show the cached results without re-running
+            st.session_state.show_results = True
+
+# Display results if the state variable is True
+if st.session_state.show_results:
+    df = st.session_state.results_df
+    if not df.empty:
+        df['category'] = df['label'].map(reverse_category_mapping)
+        st.subheader("Grouped Entities by Category", divider="violet")
+        
+        # Create tabs for each category
+        category_names = sorted(list(category_mapping.keys()))
+        category_tabs = st.tabs(category_names)
+        
+        for i, category_name in enumerate(category_names):
+            with category_tabs[i]:
+                df_category_filtered = df[df['category'] == category_name]
+                if not df_category_filtered.empty:
+                    st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
                 else:
-                    st.warning("No entities were found that occur more than once.")
-                # Download Section
-                st.divider()
-                dfa = pd.DataFrame(
-                    data={
-                        'Column Name': ['text', 'label', 'score', 'start', 'end'],
-                        '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',
-                        ]
-                    }
-                )
-                buf = io.BytesIO()
-                with zipfile.ZipFile(buf, "w") as myzip:
-                    myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
-                    myzip.writestr("Glossary of tags.csv", dfa.to_csv(index=False))
-                with stylable_container(
-                    key="download_button",
-                    css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
-                ):
-                    st.download_button(
-                        label="Download results and glossary (zip)",
-                        data=buf.getvalue(),
-                        file_name="nlpblogs_results.zip",
-                        mime="application/zip",
-                    )
-                if comet_initialized:
-                    experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
-                    experiment.end()
-            else: # If df is empty
-                st.warning("No entities were found in the provided text.")
-        end_time = time.time()
+                    st.info(f"No entities found for the '{category_name}' category.")
+        
+        with st.expander("See Glossary of tags"):
+            st.write('''
+            - **text**: ['entity extracted from your text data']
+            - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
+            - **label**: ['label (tag) assigned to a given extracted entity']
+            - **start**: ['index of the start of the corresponding entity']
+            - **end**: ['index of the end of the corresponding entity']
+            ''')
+        st.divider()
+        
+        # Tree map
+        st.subheader("Tree map", divider="violet")
+        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
+        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#E8F5E9', plot_bgcolor='#E8F5E9')
+        st.plotly_chart(fig_treemap)
+        
+        # Pie and Bar charts
+        grouped_counts = df['category'].value_counts().reset_index()
+        grouped_counts.columns = ['category', 'count']
+        col1, col2 = st.columns(2)
+        
+        with col1:
+            st.subheader("Pie chart", divider="violet")
+            fig_pie = px.pie(grouped_counts, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
+            fig_pie.update_traces(textposition='inside', textinfo='percent+label')
+            fig_pie.update_layout(
+                paper_bgcolor='#E8F5E9',
+                plot_bgcolor='#E8F5E9'
+            )
+            st.plotly_chart(fig_pie)
+            
+        with col2:
+            st.subheader("Bar chart", divider="violet")
+            fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
+            fig_bar.update_layout(
+                paper_bgcolor='#E8F5E9',
+                plot_bgcolor='#E8F5E9'
+            )
+            st.plotly_chart(fig_bar)
+            
+        # Most Frequent Entities
+        st.subheader("Most Frequent Entities", divider="gray")
+        word_counts = df['text'].value_counts().reset_index()
+        word_counts.columns = ['Entity', 'Count']
+        repeating_entities = word_counts[word_counts['Count'] > 1]
+        
+        if not repeating_entities.empty:
+            st.dataframe(repeating_entities, use_container_width=True)
+            fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
+            fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'},
+                                            paper_bgcolor='#E8F5E9',
+                                            plot_bgcolor='#E8F5E9')
+            st.plotly_chart(fig_repeating_bar)
+        else:
+            st.warning("No entities were found that occur more than once.")
+        
+        # Download Section
+        st.divider()
+        dfa = pd.DataFrame(
+            data={
+                'Column Name': ['text', 'label', 'score', 'start', 'end'],
+                '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',
+                ]
+            }
+        )
+        buf = io.BytesIO()
+        with zipfile.ZipFile(buf, "w") as myzip:
+            myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
+            myzip.writestr("Glossary of tags.csv", dfa.to_csv(index=False))
+        
+        with stylable_container(
+            key="download_button",
+            css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
+        ):
+            st.download_button(
+                label="Download results and glossary (zip)",
+                data=buf.getvalue(),
+                file_name="nlpblogs_results.zip",
+                mime="application/zip",
+            )
+        
+        if comet_initialized:
+            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
+            experiment.end()
+    else: # If df is empty
+        st.warning("No entities were found in the provided text.")
+    end_time = time.time()
     elapsed_time = end_time - start_time
     st.text("")
     st.text("")
-    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
-  
+    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.") like the excellent one?