Update src/streamlit_app.py

src/streamlit_app.py

@@ -33,21 +33,22 @@ expander.write("""**Named Entities:** This DataHarvest web app predicts nine (9)
 st.markdown("For any errors or inquiries, please contact us at [info@nlpblogs.com](mailto:info@nlpblogs.com)")
 
 with st.sidebar:
-    st.write("Use the following code to embed the DataHarvest web app on your website. Feel free to adjust the width and height values to fit your page.")
-    code = '''
-    <iframe
+   st.write("Use the following code to embed the DataHarvest web app on your website. Feel free to adjust the width and height values to fit your page.")
+   code = '''
+   <iframe
 	src="https://aiecosystem-dataharvest.hf.space"
 	frameborder="0"
 	width="850"
 	height="450"
     ></iframe>
-    '''
-    st.code(code, language="html")
-    st.text("")
-    st.text("")
-    
-    st.subheader("🚀 Ready to build your own AI Web App?", divider="violet")
-    st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", type="primary")
+ 
+   '''
+   st.code(code, language="html")
+   st.text("")
+   st.text("")
+   
+   st.subheader("🚀 Ready to build your own AI Web App?", divider="violet")
+   st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", type="primary")
 
 # --- Comet ML Setup ---
 COMET_API_KEY = os.environ.get("COMET_API_KEY")
@@ -55,39 +56,39 @@ 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.")
-    print("Warning: Comet ML environment variables are not set. Logging will be disabled.")
+   st.warning("Comet ML not initialized. Check environment variables.")
+   print("Warning: Comet ML environment variables are not set. Logging will be disabled.")
 
 # --- Label Definitions ---
 labels = ["person", "country", "city", "organization", "date", "time", "cardinal", "money", "position"]
 category_mapping = {
-    "People": ["person", "organization", "position"],
-    "Locations": ["country", "city"],
-    "Time": ["date", "time"],
-    "Numbers": ["money", "cardinal"]
+   "People": ["person", "organization", "position"],
+   "Locations": ["country", "city"],
+   "Time": ["date", "time"],
+   "Numbers": ["money", "cardinal"]
 }
 
 # --- Model Loading ---
 @st.cache_resource
 def load_ner_model():
-    """Loads the GLiNER model and caches it."""
-    try:
-        return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
-    except Exception as e:
-        st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
-        st.stop()
+   """Loads the GLiNER model and caches it."""
+   try:
+       return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
+   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()
 reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}
 
 # --- Session State Initialization ---
 if 'show_results' not in st.session_state:
-    st.session_state.show_results = False
+   st.session_state.show_results = False
 if 'last_text' not in st.session_state:
-    st.session_state.last_text = ""
+   st.session_state.last_text = ""
 if 'results_df' not in st.session_state:
-    st.session_state.results_df = pd.DataFrame()
+   st.session_state.results_df = pd.DataFrame()
 if 'elapsed_time' not in st.session_state:
-    st.session_state.elapsed_time = 0.0
+   st.session_state.elapsed_time = 0.0
 
 # --- Text Input and Clear Button ---
 word_limit = 200
@@ -96,148 +97,168 @@ word_count = len(text.split())
 st.markdown(f"**Word count:** {word_count}/{word_limit}")
 
 def clear_text():
-    """Clears the text area and hides results."""
-    st.session_state['my_text_area'] = ""
-    st.session_state.show_results = False
-    st.session_state.last_text = ""
-    st.session_state.results_df = pd.DataFrame()
-    st.session_state.elapsed_time = 0.0
+   """Clears the text area and hides results."""
+   st.session_state['my_text_area'] = ""
+   st.session_state.show_results = False
+   st.session_state.last_text = ""
+   st.session_state.results_df = pd.DataFrame()
+   st.session_state.elapsed_time = 0.0
 
 st.button("Clear text", on_click=clear_text)
 
+# --- Post-processing function to remove trailing punctuation ---
+def remove_trailing_punctuation(text_string):
+    """
+    Removes trailing punctuation from a string.
+    
+    Args:
+        text_string (str): The input string.
+        
+    Returns:
+        str: The string with trailing punctuation removed.
+    """
+    return text_string.rstrip(string.punctuation)
+
 # --- Results Section ---
 if st.button("Results"):
-    if not text.strip():
-        st.warning("Please enter some text to extract entities.")
-        st.session_state.show_results = False
-    elif word_count > word_limit:
-        st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
-        st.session_state.show_results = False
-    else:
-        # Check if the text is different from the last time
-        if text != st.session_state.last_text:
-            st.session_state.show_results = True
-            st.session_state.last_text = text
-            start_time = time.time()
-            with st.spinner("Extracting entities...", show_time=True):
-                # Pass the raw text directly to the model
-                entities = model.predict_entities(text, labels)
-                df = pd.DataFrame(entities)
-                st.session_state.results_df = df
-                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)
-                        experiment.end()
-            end_time = time.time()
-            st.session_state.elapsed_time = end_time - start_time
-            # Place the message here, so it only runs once per button click
-            st.info(f"Results processed in **{st.session_state.elapsed_time:.2f} seconds**.")
-        # If the text is the same, do nothing but keep results displayed
-        else:
-            st.session_state.show_results = True
+   if not text.strip():
+       st.warning("Please enter some text to extract entities.")
+       st.session_state.show_results = False
+   elif word_count > word_limit:
+       st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
+       st.session_state.show_results = False
+   else:
+       # Check if the text is different from the last time
+       if text != st.session_state.last_text:
+           st.session_state.show_results = True
+           st.session_state.last_text = text
+           start_time = time.time()
+           with st.spinner("Extracting entities...", show_time=True):
+               # Pass the raw text directly to the model
+               entities = model.predict_entities(text, labels)
+               df = pd.DataFrame(entities)
+               
+               # Apply post-processing to remove punctuation
+               if not df.empty:
+                   df['text'] = df['text'].apply(remove_trailing_punctuation)
+
+               st.session_state.results_df = df
+               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)
+                       experiment.end()
+               end_time = time.time()
+               st.session_state.elapsed_time = end_time - start_time
+               # Place the message here, so it only runs once per button click
+               st.info(f"Results processed in **{st.session_state.elapsed_time:.2f} seconds**.")
+       # If the text is the same, do nothing but keep results displayed
+       else:
+           st.session_state.show_results = True
 
 # Display results if the state variable is True
 if st.session_state.show_results:
-    df = st.session_state.results_df
-    if not df.empty:
-        df['category'] = df['label'].map(reverse_category_mapping)
-        st.subheader("Grouped Entities by Category", divider="violet")
-        
-        category_names = sorted(list(category_mapping.keys()))
-        category_tabs = st.tabs(category_names)
-        
-        for i, category_name in enumerate(category_names):
-            with category_tabs[i]:
-                df_category_filtered = df[df['category'] == category_name]
-                if not df_category_filtered.empty:
-                    st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
-                else:
-                    st.info(f"No entities found for the '{category_name}' category.")
-                
-        with st.expander("See Glossary of tags"):
-            st.write('''
-            - **start**: ['index of the start of the corresponding entity']
-            - **end**: ['index of the end of the corresponding entity']
-            - **text**: ['entity extracted from your text data']
-            - **label**: ['label (tag) assigned to a given extracted entity']
-            - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
-            
-            ''')
-        
-        st.divider()
-        # Tree map
-        st.subheader("Tree map", divider="violet")
-        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))
-        expander = st.expander("**Download**")
-        expander.write("""You can easily download the tree map by hovering over it. Look for the download icon that appears in the top right corner.
-        """)
-        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="violet")
-            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')
-            expander = st.expander("**Download**")
-            expander.write("""You can easily download the pie chart by hovering over it. Look for the download icon that appears in the top right corner.
-            """)
-            st.plotly_chart(fig_pie)
-            
-        with col2:
-            st.subheader("Bar chart", divider="violet")
-            fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
-            expander = st.expander("**Download**")
-            expander.write("""You can easily download the bar chart by hovering over it. Look for the download icon that appears in the top right corner.
-            """)
-            st.plotly_chart(fig_bar)
-        
-        # Most Frequent Entities
-        st.subheader("Most Frequent Entities", divider="violet")
-        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'})
-            expander = st.expander("**Download**")
-            expander.write("""You can easily download the bar chart by hovering over it. Look for the download icon that appears in the top right corner.
-            """)
-            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': ['start', 'end', 'text', 'label', 'score'],
+   df = st.session_state.results_df
+   if not df.empty:
+       df['category'] = df['label'].map(reverse_category_mapping)
+       st.subheader("Grouped Entities by Category", divider="violet")
+       
+       category_names = sorted(list(category_mapping.keys()))
+       category_tabs = st.tabs(category_names)
+       
+       for i, category_name in enumerate(category_names):
+           with category_tabs[i]:
+               df_category_filtered = df[df['category'] == category_name]
+               if not df_category_filtered.empty:
+                   st.dataframe(df_category_filtered.drop(columns=['category']), use_container_width=True)
+               else:
+                   st.info(f"No entities found for the '{category_name}' category.")
+               
+       with st.expander("See Glossary of tags"):
+           st.write('''
+           - **start**: ['index of the start of the corresponding entity']
+           - **end**: ['index of the end of the corresponding entity']
+           - **text**: ['entity extracted from your text data']
+           - **label**: ['label (tag) assigned to a given extracted entity']
+           - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
+           
+           ''')
+       
+       st.divider()
+       # Tree map
+       st.subheader("Tree map", divider="violet")
+       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))
+       expander = st.expander("**Download**")
+       expander.write("""You can easily download the tree map by hovering over it. Look for the download icon that appears in the top right corner.
+       """)
+       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="violet")
+           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')
+           expander = st.expander("**Download**")
+           expander.write("""You can easily download the pie chart by hovering over it. Look for the download icon that appears in the top right corner.
+           """)
+           st.plotly_chart(fig_pie)
+           
+       with col2:
+           st.subheader("Bar chart", divider="violet")
+           fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
+           expander = st.expander("**Download**")
+           expander.write("""You can easily download the bar chart by hovering over it. Look for the download icon that appears in the top right corner.
+           """)
+           st.plotly_chart(fig_bar)
+       
+       # Most Frequent Entities
+       st.subheader("Most Frequent Entities", divider="violet")
+       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'})
+           expander = st.expander("**Download**")
+           expander.write("""You can easily download the bar chart by hovering over it. Look for the download icon that appears in the top right corner.
+           """)
+           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': ['start', 'end', 'text', 'label', 'score'],
                                 'Description': ['index of the start of the corresponding entity', 'index of the end of the corresponding entity', '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']})
-        
-        buf = io.BytesIO()
-        with zipfile.ZipFile(buf, "w") as myzip:
-            myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
-            myzip.writestr("Most Frequent Entities.csv", repeating_entities.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: #8A2BE2; border: 1px solid black; padding: 5px; color: white; }""",
-        ):
-            st.download_button(
-                label="Download results and glossary (zip)",
-                data=buf.getvalue(),
-                file_name="nlpblogs_results.zip",
-                mime="application/zip"
-            )
-        st.text("")
-        st.text("")
-    else:
-        st.warning("No entities were found in the provided text.")
+       
+       buf = io.BytesIO()
+       with zipfile.ZipFile(buf, "w") as myzip:
+           myzip.writestr("Summary of the results.csv", df.to_csv(index=False))
+           myzip.writestr("Most Frequent Entities.csv", repeating_entities.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: #8A2BE2; border: 1px solid black; padding: 5px; color: white; }""",
+       ):
+           st.download_button(
+               label="Download results and glossary (zip)",
+               data=buf.getvalue(),
+               file_name="nlpblogs_results.zip",
+               mime="application/zip"
+           )
+       st.text("")
+       st.text("")
+   else:
+       st.warning("No entities were found in the provided text.")
+
+