app.py

import streamlit as st
import pandas as pd
import io
import plotly.express as px
import zipfile
from gliner import GLiNER
import os
from streamlit_extras.stylable_container import stylable_container

st.set_page_config(layout="wide", page_title="Named Entity Recognition App")

# --- App Header and Info ---
st.subheader("Free NER Web App", divider="red")
st.link_button("DEMO APP by nlpblogs", "https://nlpblogs.com", type="tertiary")

expander = st.expander("**Important notes on the Free NER Web App**")
expander.write('''
    **Named Entities:** This Free NER Web App predicts nine (9) labels
    grouped into three categories: **People** (person, organization, position),
    **Locations** (country, city), and **Numbers** (date, seconds, money, percent value).
    Results are presented in an easy-to-read table, visualized in an
    interactive treemap, pie chart, and bar chart, and are available for download
    along with a Glossary of tags.
    
    **How to Use:** Type or paste your text and press Ctrl + Enter. Then,
    click the 'Results' button to extract and tag entities in your text data.
    
    **Usage Limits:** Unlimited number of Result requests.
    
    **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.
    
    For any errors or inquiries, please contact us at [email protected]
    ''')

# --- Sidebar ---
with st.sidebar:
    container = st.container(border=True)
    container.write("**Named Entity Recognition (NER)** is the task of extracting and tagging entities in text data. Entities can be persons, organizations, locations, countries, products, events etc.")
    st.subheader("Related NER Web Apps", divider="red")
    st.link_button("Scandinavian JSON Entity Finder", "https://nlpblogs.com/shop/named-entity-recognition-ner/scandinavian-json-entity-finder/", 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")

if COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME:
    comet_initialized = True
else:
    comet_initialized = False
    st.warning("Comet ML not initialized. Check environment variables.")

# --- Cache the GLiNER model ---
@st.cache_resource
def load_gliner_model():
    """Caches the GLiNER model to prevent re-loading on every app rerun."""
    return GLiNER.from_pretrained("xomad/gliner-model-merge-large-v1.0")

# Load the model using the cached function
model = load_gliner_model()
# --- End Caching ---

# --- Text Input and Clear Button ---
text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", key='my_text_area')
st.write("**Input text**: ", text)

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(): # Check if the input text is empty
        st.warning("Please enter some text to extract entities.")
    else:
        with st.spinner("Extracting entities..."): # Spinner while processing
            
            # --- MODIFICATION: ADDED "seconds" to labels ---
            labels = ["person", "country", "city", "organization", "date", "seconds", "money", "percent value", "position"]
            entities = model.predict_entities(text, labels)

            df = pd.DataFrame(entities)

            # --- MODIFICATION: ADDED "seconds" to category mapping ---
            if not df.empty:
                # Create a mapping dictionary for labels to categories
                category_mapping = {
                    "person": "People",
                    "organization": "People",
                    "position": "People",
                    "country": "Locations",
                    "city": "Locations",
                    "date": "Time",
                    "seconds": "Time",
                    "money": "Numbers",
                    "percent value": "Numbers"
                }
                # Add a new 'category' column to the DataFrame
                df['category'] = df['label'].map(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()
            
            properties = {"border": "2px solid gray", "color": "blue", "font-size": "16px"}
            df_styled = df.style.set_properties(**properties)
            st.dataframe(df_styled)

            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']
                ''')

            # --- Visualizations ---
            if not df.empty: # Only plot if DataFrame is not empty
                st.subheader("Tree map", divider="red")
                # Modified treemap path to show category as the first level
                fig = px.treemap(df, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
                fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
                st.plotly_chart(fig)
                if comet_initialized:
                    experiment.log_figure(figure=fig, figure_name="entity_treemap_categories")

                col1, col2 = st.columns(2)
                with col1:
                    st.subheader("Pie Chart", divider="red")
                    # Pie chart now visualizes the distribution of categories
                    value_counts1 = df['category'].value_counts()
                    df1 = pd.DataFrame(value_counts1)
                    final_df = df1.reset_index().rename(columns={"index": "category"})
                    fig1 = px.pie(final_df, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
                    fig1.update_traces(textposition='inside', textinfo='percent+label')
                    st.plotly_chart(fig1)
                    if comet_initialized:
                        experiment.log_figure(figure=fig1, figure_name="category_pie_chart")

                with col2:
                    st.subheader("Bar Chart", divider="red")
                    # Bar chart now visualizes the distribution of categories
                    fig2 = px.bar(final_df, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
                    st.plotly_chart(fig2)
                    if comet_initialized:
                        experiment.log_figure(figure=fig2, figure_name="category_bar_chart")
            else:
                st.info("No entities found in the provided text.")

            # --- Download Buttons ---
            dfa = pd.DataFrame(
                data={
                    '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'],
                    })
            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: yellow; border: 1px solid black; padding: 5px; color: black; }""",
            ):
                st.download_button(
                    label="Download zip file",
                    data=buf.getvalue(),
                    file_name="zip file.zip",
                    mime="application/zip",
                )
            st.divider()

    end_time = time.time()
    elapsed_time = end_time - start_time
    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")