Update app.py

app.py

@@ -1,50 +1,41 @@
 import streamlit as st
-import time
+import time # Still imported, but time.sleep(5) will be removed from the main logic
 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
-from gliner import GLiNER
-
+from gliner import GLiNER # Import GLiNER
 import os
 from comet_ml import Experiment
 
-
-
-st.subheader("8-Named Entity Recognition Web App", divider = "red")
-st.link_button("DEMO APP by nlpblogs", "https://nlpblogs.com", type = "tertiary")
+# --- App Header and Info ---
+st.subheader("8-Named Entity Recognition Web App", divider="red")
+st.link_button("DEMO APP by nlpblogs", "https://nlpblogs.com", type="tertiary")
 
 expander = st.expander("**Important notes on the 8-Named Entity Recognition Web App**")
 expander.write('''
-    
-    **Named Entities:**
-    This 8-Named Entity Recognition 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 8-Named Entity Recognition 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
+    ''')
+
+# --- 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 NLP Web Apps", divider = "red")
-    st.link_button("14-Named Entity Recognition Web App", "https://nlpblogs.com/shop/named-entity-recognition-ner/14-named-entity-recognition-web-app/", type = "primary")
-
+    st.subheader("Related NLP Web Apps", divider="red")
+    st.link_button("14-Named Entity Recognition Web App", "https://nlpblogs.com/shop/named-entity-recognition-ner/14-named-entity-recognition-web-app/", 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")
@@ -55,8 +46,17 @@ else:
     comet_initialized = False
     st.warning("Comet ML not initialized. Check environment variables.")
 
+# --- Cache the GLiNER model ---
+@st.cache_resource
+def load_gliner_model():
     
+    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)
 
@@ -64,96 +64,100 @@ def clear_text():
     st.session_state['my_text_area'] = ""
 
 st.button("Clear text", on_click=clear_text)
-
 st.divider()
 
+# --- Results Section ---
 if st.button("Results"):
-    with st.spinner("Wait for it...", show_time=True):
-        time.sleep(5)
-        model = GLiNER.from_pretrained("xomad/gliner-model-merge-large-v1.0")
-        labels = ["person", "country", "city", "organization", "date", "money", "percent value", "position"]
-        entities = model.predict_entities(text, labels)
-        df = pd.DataFrame(entities)
-
-        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)
-
-        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']
-            
-            '**start**': ['index of the start of the corresponding entity']
-            
-            '**end**': ['index of the end of the corresponding entity']
-            ''')
-
-        if df is not None:
-            fig = px.treemap(df, path=[px.Constant("all"), 'text', 'label'],
-                                 values='score', color='label')
-            fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))
-            st.subheader("Tree map", divider = "red")
-            st.plotly_chart(fig)
+    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"]
+            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)
+
             if comet_initialized:
-                experiment.log_figure(figure=fig, figure_name="entity_treemap")
-
-        if df is not None:
-            value_counts1 = df['label'].value_counts()
-            df1 = pd.DataFrame(value_counts1)
-            final_df = df1.reset_index().rename(columns={"index": "label"})
-            col1, col2 = st.columns(2)
-            with col1:
-                fig1 = px.pie(final_df, values='count', names='label', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted labels')
-                fig1.update_traces(textposition='inside', textinfo='percent+label')
-                st.subheader("Pie Chart", divider = "red")
-                st.plotly_chart(fig1)
+                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)
+
+            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']
+                '**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")
+                fig = px.treemap(df, path=[px.Constant("all"), 'text', 'label'],
+                                 values='score', color='label')
+                fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+                st.plotly_chart(fig)
                 if comet_initialized:
-                    experiment.log_figure(figure=fig1, figure_name="label_pie_chart")
-            with col2:
-                fig2 = px.bar(final_df, x="count", y="label", color="label", text_auto=True, title='Occurrences of predicted labels')
-                st.subheader("Bar Chart", divider = "red")
-                st.plotly_chart(fig2)
+                    experiment.log_figure(figure=fig, figure_name="entity_treemap")
+
+                col1, col2 = st.columns(2)
+                with col1:
+                    st.subheader("Pie Chart", divider="red")
+                    value_counts1 = df['label'].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')
+                    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")
+
+                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')
+                    st.plotly_chart(fig2)
+                    if comet_initialized:
+                        experiment.log_figure(figure=fig2, figure_name="label_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'],
+                    '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",
+                )
                 if comet_initialized:
-                    experiment.log_figure(figure=fig2, figure_name="label_bar_chart")
-
-        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'],
-                '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",
-            )
-            if comet_initialized:
-                experiment.log_asset(buf.getvalue(), file_name="downloadable_results.zip")
+                    experiment.log_asset(buf.getvalue(), file_name="downloadable_results.zip")
 
-        st.divider()
-        if comet_initialized:
-            experiment.end()
+            st.divider()
+            if comet_initialized:
+                experiment.end()