Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,402 @@
-import altair as alt
-import numpy as np
-import pandas as pd
+import os
+os.environ['HF_HOME'] = '/tmp'
+import time
 import streamlit as st
+import pandas as pd
+import io
+import plotly.express as px
+import zipfile
+import json
+from cryptography.fernet import Fernet
+from streamlit_extras.stylable_container import stylable_container
+from typing import Optional
+from gliner import GLiNER
+from comet_ml import Experiment
+
+
+st.markdown(
+    """
+    <style>
+    /* Main app background and text color */
+    .stApp {
+        background-color: #E8F5E9; /* A very light green */
+        color: #1B5E20; /* Dark green for the text */
+    }
+    
+    /* Sidebar background color */
+    .css-1d36184 {
+        background-color: #A5D6A7; /* A medium light green */
+        secondary-background-color: #A5D6A7;
+    }
+    
+    /* Expander background color and header */
+    .streamlit-expanderContent, .streamlit-expanderHeader {
+        background-color: #E8F5E9;
+    }
+    
+    /* Text Area background and text color */
+    .stTextArea textarea {
+        background-color: #81C784; /* A slightly darker medium green */
+        color: #1B5E20; /* Dark green for text */
+    }
+    
+    /* Button background and text color */
+    .stButton > button {
+        background-color: #81C784;
+        color: #1B5E20;
+    }
+    
+    /* Warning box background and text color */
+    .stAlert.st-warning {
+        background-color: #66BB6A; /* A medium-dark green for the warning box */
+        color: #1B5E20;
+    }
+    
+    /* Success box background and text color */
+    .stAlert.st-success {
+        background-color: #66BB6A; /* A medium-dark green for the success box */
+        color: #1B5E20;
+    }
+    </style>
+    """,
+    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")
+st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
+
+expander = st.expander("**Important notes**")
+expander.write("""**Named Entities:** This PiiGuard web app predicts eight (8) labels:"Location", "Organization", "Product_or_Good", "Date", "Quantity", "Transportation_Mode", "Person", "Document_or_Form_ID"
+
+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.      
+
+**Usage Limits:** You can request results unlimited times for one (1) month.      
+
+**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 = '''
+    <iframe
+	src="https://aiecosystem-chainsense.hf.space"
+	frameborder="0"
+	width="850"
+	height="450"
+    ></iframe>
+   
+    '''
+    st.code(code, language="html")
+    st.text("")
+    st.text("")
+    st.divider()
+    st.subheader("🚀 Ready to build your own AI Web App?", divider="violet")
+    st.link_button("AI Web App Builder", " https://nlpblogs.com/custom-web-app-development/", type="primary")
+
+# --- Comet ML Setup ---
+COMET_API_KEY = os.environ.get("COMET_API_KEY")
+COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
+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": [
+    "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_resource
+def load_ner_model():
+    """Loads the GLiNER model and caches it."""
+    try:
+        return GLiNER.from_pretrained("E3-JSI/gliner-multi-pii-domains-v1", 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}
+
+# --- Text Input and Clear Button ---
+text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area')
+
+def clear_text():
+    """Clears the text area."""
+    st.session_state['my_text_area'] = ""
+
+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.")
+    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)
+                
+                
+
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+                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="violet")
+                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**.")