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>
@@ -58,13 +59,14 @@ st.markdown(
     }
     </style>
     """,
-    unsafe_allow_html=True)
+    unsafe_allow_html=True
+)
 # --- Page Configuration and UI Elements ---
 st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 st.subheader("Legal Lens", divider="grey")
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
 expander = st.expander("**Important notes**")
-expander.write("""**Named Entities:** This Legal Lens web app predicts twenty-eight (28) labels: "Plaintiff", "Defendant", "Appellant", "Appellee", "Debtor", "Creditor", "Signer", "Witness", "Courts", "Judges", "Lawyers", "Attorneys", "Statutes", "Laws", "Provisions", "Case_citations", "Legal_documents", "Effective_dates", "Execution_dates", "Expiration_dates", "Money", "Amounts", "Contract_terms", "Case_number", "Witnesses", "Crimes", "Offenses", "Victims"
+expander.write("""**Named Entities:** This Legal Lens web app predicts twenty-eight (28) labels: "Plaintiff", "Defendant", "Appellant", "Appellee", "Debtor", "Creditor", "Signer", "Witness", "Courts", "Judges", "Lawyers", "Attorneys", "Statutes", "Laws", "Provisions", "Case_citations", "Legal_documents", "Effective_dates", "Execution_dates", "Expiration_dates", "Money", "Amounts", "Contract_terms", "Case_number", "Crimes", "Offenses", "Victims"
 
 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.
 
@@ -103,7 +105,7 @@ comet_initialized = bool(COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAM
 if not comet_initialized:
     st.warning("Comet ML not initialized. Check environment variables.")
 # --- Label Definitions ---
-labels = ["Plaintiff","Defendant","Appellant","Appellee","Debtor","Creditor","Signer","Witness","Courts","Judges","Lawyers","Attorneys","Statutes","Laws","Provisions","Case_citations","Legal_documents","Effective_dates","Execution_dates","Expiration_dates","Money","Amounts","Contract_terms","Case_number","Witnesses","Crimes","Offenses","Victims"]
+labels = ["Plaintiff","Defendant","Appellant","Appellee","Debtor","Creditor","Signer","Witness","Courts","Judges","Lawyers","Attorneys","Statutes","Laws","Provisions","Case_citations","Legal_documents","Effective_dates","Execution_dates","Expiration_dates","Money","Amounts","Contract_terms","Case_number","Crimes","Offenses","Victims"]
 category_mapping = {
     "Parties": [
         "Plaintiff",
@@ -142,14 +144,12 @@ category_mapping = {
     ],
     "Court Judgments": [
         "Case_number",
-        "Witnesses",
     ],
     "Criminal Law": [
         "Crimes",
         "Offenses",
         "Victims"
-    ]
-}
+    ]}
 # --- Model Loading ---
 @st.cache_resource
 def load_ner_model():
@@ -192,7 +192,7 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df)
-                st.subheader("Grouped Entities by Category", divider = "grey")
+                st.subheader("Grouped Entities by Category", divider="grey")
                 # Create tabs for each category
                 category_names = sorted(list(category_mapping.keys()))
                 category_tabs = st.tabs(category_names)
@@ -213,7 +213,7 @@ if st.button("Results"):
                     ''')
                 st.divider()
                 # Tree map
-                st.subheader("Tree map", divider = "grey")
+                st.subheader("Tree map", divider="grey")
                 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), paper_bgcolor='#F5F5F5', plot_bgcolor='#F5F5F5')
                 st.plotly_chart(fig_treemap)
@@ -222,7 +222,7 @@ if st.button("Results"):
                 grouped_counts.columns = ['category', 'count']
                 col1, col2 = st.columns(2)
                 with col1:
-                    st.subheader("Pie chart", divider = "grey")
+                    st.subheader("Pie chart", divider="grey")
                     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(
@@ -231,9 +231,9 @@ if st.button("Results"):
                     )
                     st.plotly_chart(fig_pie)
                 with col2:
-                    st.subheader("Bar chart", divider = "grey")
+                    st.subheader("Bar chart", divider="grey")
                     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='#F5F5F5',
                         plot_bgcolor='#F5F5F5'
                     )
@@ -272,7 +272,7 @@ if st.button("Results"):
                     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; }""",
+                    css_styles="""button { background-color: #8C8C8C; border: 1px solid black; padding: 5px; color: white; }""",
                 ):
                     st.download_button(
                         label="Download results and glossary (zip)",
@@ -283,10 +283,12 @@ if st.button("Results"):
                 if comet_initialized:
                     experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
                     experiment.end()
+                
+                # Correct placement of timing information
+                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.")