Update src/streamlit_app.py

src/streamlit_app.py

@@ -12,9 +12,6 @@ from streamlit_extras.stylable_container import stylable_container
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
-
-
-
 st.markdown(
     """
     <style>
@@ -23,24 +20,20 @@ st.markdown(
         background-color: #F5F5F5; /* A very light grey */
         color: #333333; /* Dark grey for text for good contrast */
     }
-
     /* Sidebar background */
     .css-1d36184, .css-1d36184, .st-ck {
         background-color: #D3D3D3; /* Light grey for the sidebar */
     }
-
     /* Expander header and content background */
     .streamlit-expanderHeader, .streamlit-expanderContent {
         background-color: #F5F5F5;
     }
-
     /* Text Area background and text color */
     .stTextArea textarea {
         background-color: #E6E6E6; /* Slightly darker grey for input fields */
         color: #000000;
         border: 1px solid #B0B0B0; /* Add a subtle border */
     }
-
     /* Button styling */
     .stButton > button {
         background-color: #B0B0B0; /* A medium grey for the button */
@@ -52,7 +45,6 @@ st.markdown(
     .stButton > button:hover {
         background-color: #8C8C8C; /* Darker grey on hover */
     }
-
     /* Alert boxes */
     .stAlert {
         color: #000000;
@@ -66,42 +58,18 @@ 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"
-
-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.      
-
-**Usage Limits:** You can request results unlimited times for one (1) month.      
-
-**Supported Languages:** English 
-
-**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""")
-
+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"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.      **Usage Limits:** You can request results unlimited times for one (1) month.      **Supported Languages:** English **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 Legal Lens web app on your website. Feel free to adjust the width and height values to fit your page.")
     code = '''
-   <iframe
-	src="https://aiecosystem-legal-lens.hf.space"
-	frameborder="0"
-	width="850"
-	height="450"
+   <iframe	src="https://aiecosystem-legal-lens.hf.space"	frameborder="0"	width="850"	height="450"
     ></iframe>
-
     '''
     st.code(code, language="html")
     st.text("")
@@ -109,57 +77,19 @@ with st.sidebar:
     st.divider()
     st.subheader("🚀 Ready to build your own AI Web App?", divider="grey")
     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")
 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 = [
-"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","Witnesses","Crimes","Offenses","Victims"]
 category_mapping = {
     "Parties": [
         "Plaintiff",
-        
         "Defendant",
-        
         "Appellant",
         "Appellee",
         "Debtor",
@@ -167,67 +97,41 @@ category_mapping = {
         "Signer",
         "Witness"
     ],
-    
     "Judicial & Governmental Bodies": [
         "Courts",
         "Judges",
         "Lawyers",
         "Attorneys"
-
     ],
-    
     "Legal Instruments & Concepts": [
         "Statutes",
         "Laws",
         "Provisions",
         "Case_citations",
         "Legal_documents"
-
     ],
-    
     "Dates & Timeframes": [
         "Effective_dates",
         "Execution_dates",
         "Expiration_dates"
-
-
     ],
-    
     "Financial & Monetary Entities": [
         "Money",
         "Amounts"
-        
     ],
-    
     "Contracts": [
         "Contract_terms"
-        
     ],
-    
     "Court Judgments": [
         "Case_number",
         "Witnesses",
-       
     ],
-
-
-
-     "Criminal Law": [
+    "Criminal Law": [
         "Crimes",
         "Offenses",
         "Victims"
-       
     ]
-
-    
 }
-
-
-
-
-
-
-
 # --- Model Loading ---
 @st.cache_resource
 def load_ner_model():
@@ -238,30 +142,28 @@ def load_ner_model():
         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
 reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}
-
 # --- Text Input and Clear Button ---
-text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area')
-
+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:
         with st.spinner("Extracting entities...", show_time=True):
             entities = model.predict_entities(text, labels)
             df = pd.DataFrame(entities)
-
             if not df.empty:
                 df['category'] = df['label'].map(reverse_category_mapping)
                 if comet_initialized:
@@ -272,13 +174,10 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df)
-                
                 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)
-                
                 for i, category_name in enumerate(category_names):
                     with category_tabs[i]:
                         df_category_filtered = df[df['category'] == category_name]
@@ -286,9 +185,6 @@ if st.button("Results"):
                             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('''
                     - **text**: ['entity extracted from your text data']
@@ -298,18 +194,15 @@ if st.button("Results"):
                     - **end**: ['index of the end of the corresponding entity']
                     ''')
                 st.divider()
-                                
                 # Tree map
                 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)
-                
                 # 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 = "grey")
                     fig_pie = px.pie(grouped_counts, values='count', names='category', hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
@@ -319,10 +212,6 @@ if st.button("Results"):
                         plot_bgcolor='#F5F5F5'
                     )
                     st.plotly_chart(fig_pie)
-                
-                
-
-
                 with col2:
                     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')
@@ -331,7 +220,6 @@ if st.button("Results"):
                         plot_bgcolor='#F5F5F5'
                     )
                     st.plotly_chart(fig_bar)
-                
                 # Most Frequent Entities
                 st.subheader("Most Frequent Entities", divider="grey")
                 word_counts = df['text'].value_counts().reset_index()
@@ -346,10 +234,8 @@ 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'],
@@ -359,7 +245,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',
-                            
                         ]
                     }
                 )
@@ -367,7 +252,6 @@ 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; }""",
@@ -378,14 +262,12 @@ 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("")