src/streamlit_app.py

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: white;
        color: black;
    }
    /* Sidebar background color */
    .css-1d36184 {
        background-color: #ADD8E6;
        secondary-background-color: #ADD8E6;
    }
        
    /* Expander background color */
    .streamlit-expanderContent {
        background-color: white;
    }
    /* Expander header background color */
    .streamlit-expanderHeader {
        background-color: white;
    }
    /* Text Area background and text color */
    .stTextArea textarea {
        background-color: lavender;
        color: black;
    }
    /* Button background and text color */
    .stButton > button {
        background-color: lavender;
        color: black;
    }
    /* Warning box background and text color */
    .stAlert.st-warning {
        background-color: #lavender;
        color: black;
    }
    /* Success box background and text color */
    .stAlert.st-success {
        background-color: #lavender;
        color: black;
    }
    </style>
    """,
    unsafe_allow_html=True
)






# --- Page Configuration and UI Elements ---
st.set_page_config(layout="wide", page_title="Named Entity Recognition App")

st.subheader("Public Service", divider="gray")
st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")

expander = st.expander("**Important notes on the ProductTag**")
expander.write("""
   **Named Entities:** This ProductTag predicts twenty-four (24) labels: "Product", "Service", "Organization", "Company", "Currency", "City", "Country", "Region", "Market", "Store", "Shop", "Customer_segment", "Demographics", "Target_market", "Market_segment", "Fiscal_period", "Timeframe", "Date", "Campaign", "Advertisement", "Event", "Media_platform", "Media_channel", "Social_media_platform"
   
   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) week.
   
   **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 [email protected]
""")

with st.sidebar:
    st.subheader("Build your own NER Web App in a minute without writing a single line of code.", divider="gray")
    st.link_button("NER File Builder", "https://nlpblogs.com/shop/named-entity-recognition-ner/ner-file-builder/", type="primary")

    st.text("")
    st.text("")

    st.write("Use the following code to embed the ProductTag web app on your website. Feel free to adjust the width and height values to fit your page.")
    code = '''
    <iframe
    src="https://aiecosystem-producttag1.hf.space"
    frameborder="0"
    width="850"
    height="450"
    ></iframe>
    '''
    st.code(code, language="html")

# --- 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",
        "Date of birth",
        "Blood type",
        "Digital signature",
        "Social media handle",
        "Username",
        "Birth certificate number",
        "Address",
        "Phone number",
        "Mobile phone number",
        "Landline phone number",
        "Email",
        "Fax number",
        "Postal code",
        "Credit card number",
        "Credit card expiration date",
        "CVV",
        "CVC",
        "Bank account number",
        "IBAN",
        "Transaction number",
        "Credit card brand",
        "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",
        "Medication",
        "Medical condition",
        "Health insurance ID number",
        "Health insurance number",
        "National health insurance number",
        "Flight number",
        "Reservation number",
        "Train ticket number",
        "Vehicle registration number",
        "Organization",
        "Insurance company",
        "IP address",
        "Serial number",
        "Insurance number",
    
]

category_mapping = {
    "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",
    ]
}






# --- Model Loading ---
@st.cache_resource
def load_ner_model():
    """Loads the GLiNER model and caches it."""
    try:
        return GLiNER.from_pretrained("urchade/gliner_multi_pii-v1", nested_ner=True, num_gen_sequences=2, gen_constraints= labels, threshold = 0.70)
    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)
st.divider()

# --- 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("Extracted Entities", divider = "gray")
                st.dataframe(df.style.set_properties(**{"border": "2px solid gray", "color": "blue", "font-size": "16px"}))

                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']
                    - **category**: ['the high-level category for the label']
                    - **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 = "gray")
                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))
               

                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 = "gray")
                    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')
                    st.plotly_chart(fig_pie)

                with col2:
                    st.subheader("Bar chart", divider = "gray")
                    fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True,
                                    title='Occurrences of predicted categories')
                    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'})
                    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', '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 the entity belongs to',
                        ]
                    }
                )

                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="markettag_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**.")