Update app.py

app.py

@@ -1,13 +1,11 @@
 import streamlit as st
-import time # Still imported, but time.sleep(5) will be removed from the main logic
 import pandas as pd
 import io
-from streamlit_extras.stylable_container import stylable_container
 import plotly.express as px
 import zipfile
-from gliner import GLiNER # Import GLiNER
+from gliner import GLiNER
 import os
-from comet_ml import Experiment
+from streamlit_extras.stylable_container import stylable_container
 
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 
@@ -17,17 +15,26 @@ 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 eight (8) labels (“person”, “country”, “city”, “organization”, “date”, “money”, “percent value”, “position”). 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:** 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 info@nlpblogs.com
+    **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 ---
@@ -51,7 +58,7 @@ else:
 # --- 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
@@ -63,6 +70,7 @@ text = st.text_area("Type or paste your text below, and then press Ctrl + Enter"
 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)
@@ -70,18 +78,35 @@ 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
-            # No need for time.sleep(5) here, as the model is already loaded
-            labels = ["person", "country", "city", "organization", "date", "money", "percent value", "position"]
+            
+            # --- MODIFICATION: ADDED "seconds" to labels ---
+            labels = ["person", "country", "city", "organization", "date", "seconds", "money", "percent value", "position"]
             entities = model.predict_entities(text, labels)
 
-            # Ensure entities is a list of dictionaries for DataFrame creation
-            # If no entities are found, 'entities' might be an empty list, which is fine for pd.DataFrame
             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,
@@ -90,7 +115,8 @@ if st.button("Results"):
                 )
                 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)
@@ -103,6 +129,8 @@ if st.button("Results"):
                 
                 '**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']
@@ -111,31 +139,33 @@ if st.button("Results"):
             # --- Visualizations ---
             if not df.empty: # Only plot if DataFrame is not empty
                 st.subheader("Tree map", divider="red")
-                fig = px.treemap(df, path=[px.Constant("all"), 'text', 'label'],
-                                 values='score', color='label')
+                # 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")
+                    experiment.log_figure(figure=fig, figure_name="entity_treemap_categories")
 
                 col1, col2 = st.columns(2)
                 with col1:
                     st.subheader("Pie Chart", divider="red")
-                    value_counts1 = df['label'].value_counts()
+                    # 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": "label"})
-                    fig1 = px.pie(final_df, values='count', names='label', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted labels')
+                    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="label_pie_chart")
+                        experiment.log_figure(figure=fig1, figure_name="category_pie_chart")
 
                 with col2:
                     st.subheader("Bar Chart", divider="red")
-                    fig2 = px.bar(final_df, x="count", y="label", color="label", text_auto=True, title='Occurrences of predicted labels')
+                    # 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="label_bar_chart")
+                        experiment.log_figure(figure=fig2, figure_name="category_bar_chart")
             else:
                 st.info("No entities found in the provided text.")
 
@@ -143,6 +173,7 @@ if st.button("Results"):
             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'],
                     })
@@ -161,9 +192,11 @@ if st.button("Results"):
                     file_name="zip file.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**.")
+
+
+