Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,6 +12,7 @@ from streamlit_extras.stylable_container import stylable_container
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
+
 st.markdown(
     """
     <style>
@@ -20,36 +21,36 @@ st.markdown(
         background-color: #E8F5E9; /* A very light green */
         color: #1B5E20; /* Dark green for the text */
     }
-    
+
     /* Sidebar background color */
     .css-1d36184 {
         background-color: #A5D6A7; /* A medium light green */
         secondary-background-color: #A5D6A7;
     }
-    
+
     /* Expander background color and header */
     .streamlit-expanderContent, .streamlit-expanderHeader {
         background-color: #E8F5E9;
     }
-    
+
     /* Text Area background and text color */
     .stTextArea textarea {
         background-color: #81C784; /* A slightly darker medium green */
         color: #1B5E20; /* Dark green for text */
     }
-    
+
     /* Button background and text color */
     .stButton > button {
         background-color: #81C784;
         color: #1B5E20;
     }
-    
+
     /* Warning box background and text color */
     .stAlert.st-warning {
         background-color: #66BB6A; /* A medium-dark green for the warning box */
         color: #1B5E20;
     }
-    
+
     /* Success box background and text color */
     .stAlert.st-success {
         background-color: #66BB6A; /* A medium-dark green for the success box */
@@ -58,6 +59,7 @@ st.markdown(
     </style>
     """,
     unsafe_allow_html=True)
+
 # --- Page Configuration and UI Elements ---
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 st.subheader("EntityFinance", divider="violet")
@@ -77,7 +79,6 @@ Results are presented in easy-to-read tables, visualized in an interactive tree
 
 For any errors or inquiries, please contact us at [email protected]""")
 
-
 with st.sidebar:
     st.write("Use the following code to embed the EntityFinance web app on your website. Feel free to adjust the width and height values to fit your page.")
     code = '''
@@ -95,6 +96,7 @@ with st.sidebar:
     st.divider()
     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")
 COMET_WORKSPACE = os.environ.get("COMET_WORKSPACE")
@@ -102,6 +104,7 @@ 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.")
+
 # --- Label Definitions ---
 labels = [
     "Monetary_value",
@@ -118,20 +121,16 @@ labels = [
     "Location",
     "Date",
     "Time"]
+
 # Corrected mapping dictionary
 # Create a mapping dictionary for labels to categories
 category_mapping = {
-    "People & Groups": [ "Person",
-      "Organization",
-      "Regulatory_entity"],
-    "Financial & Transactional": [ "Monetary_value",
-      "Financial_instrument",
-      "Company_identifier",
-      "Financial_event",
-      "Financial_metric", "Product", "Service"],
+    "People & Groups": [ "Person", "Organization", "Regulatory_entity"],
+    "Financial & Transactional": [ "Monetary_value", "Financial_instrument", "Company_identifier", "Financial_event", "Financial_metric", "Product", "Service"],
     "Temporal": ["Date", "Time"],
     "Locations": ["Location"],
     "Documents & Context": ["Financial_document"]}
+
 # --- Model Loading ---
 @st.cache_resource
 def load_ner_model():
@@ -144,23 +143,27 @@ def load_ner_model():
 model = load_ner_model()
 # Flatten the mapping to a single dictionary
 reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}
+
 # --- Text Input and Clear Button ---
 word_limit = 200
 text = st.text_area(f"Type or paste your text below (max {word_limit} words), and then press Ctrl + Enter", height=250, key='my_text_area')
 word_count = len(text.split())
 st.markdown(f"**Word count:** {word_count}/{word_limit}")
+
 def clear_text():
     """Clears the text area."""
     st.session_state['my_text_area'] = ""
+
 st.button("Clear text", on_click=clear_text)
+
 # --- Results Section ---
 if st.button("Results"):
-    start_time = time.time()
     if not text.strip():
         st.warning("Please enter some text to extract entities.")
     elif word_count > word_limit:
         st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.")
     else:
+        start_time = time.time()
         with st.spinner("Extracting entities...", show_time=True):
             entities = model.predict_entities(text, labels)
             df = pd.DataFrame(entities)
@@ -215,7 +218,7 @@ if st.button("Results"):
                 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
+                    fig_bar.update_layout(
                         paper_bgcolor='#E8F5E9',
                         plot_bgcolor='#E8F5E9'
                     )
@@ -265,10 +268,12 @@ if st.button("Results"):
                 if comet_initialized:
                     experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
                     experiment.end()
+                
+                # Correctly placed time calculation
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                st.text("")
+                st.text("")
+                st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
             else: # If df is empty
-                st.warning("No entities were found in the provided text.")
-        end_time = time.time()
-    elapsed_time = end_time - start_time
-    st.text("")
-    st.text("")
-    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
+                st.warning("No entities were found in the provided text.")