Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,17 +7,13 @@ import io
 import plotly.express as px
 import zipfile
 import json
+import string
 from cryptography.fernet import Fernet
 from streamlit_extras.stylable_container import stylable_container
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
 
-
- 
-   
-
-
 # --- Page Configuration and UI Elements ---
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 st.subheader("DataHarvest", divider="violet")
@@ -27,15 +23,13 @@ st.markdown(':rainbow[**Supported Languages: English**]')
 expander = st.expander("**Important notes**")
 expander.write("""**Named Entities:** This DataHarvest web app predicts nine (9) labels: "person", "country", "city", "organization", "date", "time", "money", "percent", "position"
 
-Results are presented in easy-to-read tables, 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 into the text area below, then press Ctrl + Enter. Click the 'Results' button to extract and tag entities in your text data.      
+Results are presented in easy-to-read tables, visualized in an interactive tree map, pie chart and bar chart, and are available for download along with a Glossary of tags.
 
-**Usage Limits:** You can request results unlimited times for one (1) month.      
+**How to Use:** Type or paste your text into the text area below, then press Ctrl + Enter. Click the 'Results' button to extract and tag entities in your text data.
 
-**Technical issues:** If your connection times out, please refresh the page or reopen the app's URL. 
+**Usage Limits:** You can request results unlimited times for one (1) month.
 
-For any errors or inquiries, please contact us at info@nlpblogs.com""")
+**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:
     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.")
@@ -46,7 +40,7 @@ with st.sidebar:
 	width="850"
 	height="450"
     ></iframe>
- 
+
     '''
     st.code(code, language="html")
     st.text("")
@@ -60,7 +54,6 @@ 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 = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME)
-
 if not comet_initialized:
     st.warning("Comet ML not initialized. Check environment variables.")
 
@@ -68,7 +61,6 @@ if not comet_initialized:
 labels = ["person", "country", "city", "organization", "date", "time", "money", "percent", "position"]
 
 # Corrected mapping dictionary
-
 # Create a mapping dictionary for labels to categories
 category_mapping = {
     "People": ["person", "organization", "position"],
@@ -77,9 +69,6 @@ category_mapping = {
     "Numbers": ["money", "percent"]
 }
 
-
-
-
 # --- Model Loading ---
 @st.cache_resource
 def load_ner_model():
@@ -89,6 +78,7 @@ def load_ner_model():
     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()
 
 # Flatten the mapping to a single dictionary
@@ -101,8 +91,12 @@ def clear_text():
     """Clears the text area."""
     st.session_state['my_text_area'] = ""
 
-st.button("Clear text", on_click=clear_text)
+def remove_punctuation(text):
+    """Removes punctuation from a string."""
+    translator = str.maketrans('', '', string.punctuation)
+    return text.translate(translator)
 
+st.button("Clear text", on_click=clear_text)
 
 # --- Results Section ---
 if st.button("Results"):
@@ -110,10 +104,12 @@ if st.button("Results"):
     if not text.strip():
         st.warning("Please enter some text to extract entities.")
     else:
+        # Call the new function to remove punctuation from the input text
+        cleaned_text = remove_punctuation(text)
         with st.spinner("Extracting entities...", show_time=True):
-            entities = model.predict_entities(text, labels)
+            # Use the cleaned text for prediction
+            entities = model.predict_entities(cleaned_text, labels)
             df = pd.DataFrame(entities)
-
             if not df.empty:
                 df['category'] = df['label'].map(reverse_category_mapping)
                 if comet_initialized:
@@ -124,13 +120,13 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df)
-                
+
                 st.subheader("Grouped Entities by Category", divider = "violet")
-                
+
                 # Create tabs for each category
                 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]
@@ -139,8 +135,6 @@ if st.button("Results"):
                         else:
                             st.info(f"No entities found for the '{category_name}' category.")
 
-                
-
                 with st.expander("See Glossary of tags"):
                     st.write('''
                     - **text**: ['entity extracted from your text data']
@@ -150,38 +144,33 @@ if st.button("Results"):
                     - **end**: ['index of the end of the corresponding 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))
                 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')
                     fig_pie.update_layout(
-                        
-                    )
+                                            )
                     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')
                     fig_bar.update_layout(  # Changed from fig_pie to fig_bar
-                        
-                    )
+                                            )
                     st.plotly_chart(fig_bar)
-                
+
                 # Most Frequent Entities
                 st.subheader("Most Frequent Entities", divider="violet")
                 word_counts = df['text'].value_counts().reset_index()
@@ -195,10 +184,10 @@ if st.button("Results"):
                     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': ['text', 'label', 'score', 'start', 'end'],
@@ -208,7 +197,6 @@ if st.button("Results"):
                             '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',
-                            
                         ]
                     }
                 )
@@ -216,7 +204,7 @@ if st.button("Results"):
                 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: red; border: 1px solid black; padding: 5px; color: white; }""",
@@ -227,14 +215,13 @@ if st.button("Results"):
                         file_name="nlpblogs_results.zip",
                         mime="application/zip",
                     )
-                                
+
                 if comet_initialized:
                     experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
                     experiment.end()
             else: # If df is empty
                 st.warning("No entities were found in the provided text.")
-    
-    end_time = time.time()
+        end_time = time.time()
     elapsed_time = end_time - start_time
     st.text("")
     st.text("")