app.py

import time
import streamlit as st
import pandas as pd
import io
from transformers import pipeline
from streamlit_extras.stylable_container import stylable_container
import plotly.express as px
import zipfile
import os
from comet_ml import Experiment
import re
import numpy as np
from bs4 import BeautifulSoup 

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



# --- Configuration ---
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 = False
if COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME:
    comet_initialized = True

# --- Initialize session state ---
if 'file_upload_attempts' not in st.session_state:
    st.session_state['file_upload_attempts'] = 0

max_attempts = 10

# --- Helper function for model loading ---
@st.cache_resource
def load_ner_model():
    """Loads the pre-trained NER model and caches it."""
    return pipeline("token-classification", model="dslim/bert-base-NER", aggregation_strategy="max")

# --- UI Elements ---
st.subheader("English HTML Entity Finder", divider="rainbow")
st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")

expander = st.expander("**Important notes on the English HTML Entity Finder**")
expander.write('''
        **Named Entities:** This English HTML Entity Finder predicts four (4) labels (“PER: person”, “LOC: location”, “ORG: organization”, “MISC: miscellaneous”). 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.
        
        **How to Use:** Upload your html 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.
        
        **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 info@nlpblogs.com
    ''')

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 NLP Web Apps", divider="rainbow")
    st.link_button("Italian URL & TXT Entity Finder", "https://nlpblogs.com/shop/named-entity-recognition-ner/monolingual-ner-web-apps/italian-url-txt-entity-finder/", type="primary")






uploaded_file = st.file_uploader("Choose an HTML file", type="html")
text = None
df = None


if uploaded_file is not None:
    # Read the content of the uploaded HTML file
    html_content = uploaded_file.read().decode("utf-8")

    # Display the HTML content using components.html
    
    st.success("File uploaded successfully. Due to security protocols, the file content is hidden.")

    # --- NEW: Extract plain text from HTML ---
    soup = BeautifulSoup(html_content, 'html.parser')
    text = soup.get_text() # This extracts all visible text

   

   

st.divider()

# --- Results Button and Processing Logic ---
if st.button("Results"):
    start_time = time.time()
    if not comet_initialized:
        st.warning("Comet ML not initialized. Check environment variables if you wish to log data.")

    if st.session_state['file_upload_attempts'] >= max_attempts:
        st.error(f"You have requested results {max_attempts} times. You have reached your daily request limit.")
        st.stop()

    if text is None:
        st.warning("Please upload a supported file (.pdf or .docx) before requesting results.")
        st.stop()

    st.session_state['file_upload_attempts'] += 1

    with st.spinner("Analyzing text...", show_time=True):
        
        model = load_ner_model()
        text_entities = model(text)
        df = pd.DataFrame(text_entities)

        
        
        if 'word' in df.columns:
            st.write("Data type of 'word' column:", df['word'].dtype)


        
        
        pattern = r'[^\w\s]'
        if 'word' in df.columns: # Add this check
            df['word'] = df['word'].replace(pattern, '', regex=True)
        else:
            st.error("The 'word' column does not exist in the DataFrame. Cannot perform cleaning.")
            st.stop() # Stop execution if the column is missing
        df = df.replace('', 'Unknown').dropna()

        

        if df.empty:
            st.warning("No entities were extracted from the uploaded text.")
            st.stop()

        if comet_initialized:
            experiment = Experiment(
                api_key=COMET_API_KEY,
                workspace=COMET_WORKSPACE,
                project_name=COMET_PROJECT_NAME,
            )
            experiment.log_parameter("input_text_length", len(text))
            experiment.log_table("predicted_entities", df)

        # --- Display Results ---
        properties = {"border": "2px solid gray", "color": "blue", "font-size": "16px"}
        df_styled = df.style.set_properties(**properties)
        st.dataframe(df_styled, use_container_width=True)

        with st.expander("See Glossary of tags"):
            st.write('''
            '**word**': ['entity extracted from your text data']
            
            '**score**': ['accuracy score; how accurately a tag has been assigned to a given entity']
            
            '**entity_group**': ['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']
            ''')

        
        entity_groups = {"PER": "person",
                         "LOC": "location",
                         "ORG": "organization",
                         "MISC": "miscellaneous",
                         
                         
                        }



        st.subheader("Grouped entities", divider = "blue")
        
       
        # Convert entity_groups dictionary to a list of (key, title) tuples
        entity_items = list(entity_groups.items())
        # Define how many tabs per row
        tabs_per_row = 5
        # Loop through the entity items in chunks
        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.")
                        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.divider()


        
            
        
        
        # --- Visualizations ---
        st.subheader("Tree map", divider="rainbow")
        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'word', 'entity_group'],
                                 values='score', color='entity_group')
        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25))
        st.plotly_chart(fig_treemap)
        if comet_initialized:
            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")

        value_counts1 = df['entity_group'].value_counts()
        final_df_counts = value_counts1.reset_index().rename(columns={"index": "entity_group"})

        col1, col2 = st.columns(2)
        with col1:
            st.subheader("Pie Chart", divider="rainbow")
            fig_pie = px.pie(final_df_counts, values='count', names='entity_group', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted labels')
            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="rainbow")
            fig_bar = px.bar(final_df_counts, x="count", y="entity_group", color="entity_group", text_auto=True, title='Occurrences of predicted labels')
            st.plotly_chart(fig_bar)
            if comet_initialized:
                experiment.log_figure(figure=fig_bar, figure_name="label_bar_chart")

        # --- Downloadable Content ---
        dfa = pd.DataFrame(
                data={
                    'Column Name': ['word', 'entity_group','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: yellow; border: 1px solid black; padding: 5px; color: black; }""",
        ):
            st.download_button(
                label="Download zip file",
                data=buf.getvalue(),
                file_name="nlpblogs_ner_results.zip",
                mime="application/zip",
            )
            if comet_initialized:
                experiment.log_asset(buf.getvalue(), file_name="downloadable_results.zip")

        st.divider()
        if comet_initialized:
            experiment.end()
    end_time = time.time()
    elapsed_time = end_time - start_time
    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
st.write(f"Number of times you requested results: **{st.session_state['file_upload_attempts']}/{max_attempts}**")