Update src/streamlit_app.py

src/streamlit_app.py

@@ -13,6 +13,7 @@ from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
 
+# --- CSS Styling for the App ---
 st.markdown(
     """
     <style>
@@ -21,31 +22,37 @@ st.markdown(
         background-color: #E8F5E9; /* A very light green */
         color: #1B5E20; /* Dark green for the text */
     }
-        /* Sidebar background color */
+
+    /* Sidebar background color */
     .css-1d36184 {
         background-color: #A5D6A7; /* A medium light green */
         secondary-background-color: #A5D6A7;
     }
-        /* Expander background color and header */
+
+    /* Expander background color and header */
     .streamlit-expanderContent, .streamlit-expanderHeader {
         background-color: #E8F5E9;
     }
-        /* Text Area background and text color */
+
+    /* 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 */
+
+    /* Button background and text color */
     .stButton > button {
         background-color: #81C784;
         color: #1B5E20;
     }
-        /* Warning box background and text color */
+
+    /* 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 */
+
+    /* Success box background and text color */
     .stAlert.st-success {
         background-color: #66BB6A; /* A medium-dark green for the success box */
         color: #1B5E20;
@@ -102,123 +109,47 @@ if not comet_initialized:
 
 # --- Label Definitions ---
 labels = [
-    "person",
-    "organization",
-    "social_media_handle",
-    "username",
-    "insurance_company",
-    "phone_number",
-    "email",
-    "email_address",
-    "mobile_phone_number",
-    "landline_phone_number",
-    "fax_number",
-    "credit_card_number",
-    "credit_card_expiration_date",
-    "credit_card_brand",
-    "cvv",
-    "cvc",
-    "bank_account_number",
-    "iban",
-    "transaction_number",
-    "cpf",
-    "cnpj",
-    "passport_number",
-    "passport_expiration_date",
-    "driver's_license_number",
-    "tax_identification_number",
-    "identity_card_number",
-    "national_id_number",
-    "identity_document_number",
-    "birth_certificate_number",
-    "social_security_number",
-    "health_insurance_id_number",
-    "health_insurance_number",
-    "national_health_insurance_number",
-    "student_id_number",
-    "registration_number",
-    "insurance_number",
-    "serial_number",
-    "visa_number",
-    "reservation_number",
-    "train_ticket_number",
-    "medication",
-    "medical_condition",
-    "blood_type",
-    "date_of_birth",
-    "address",
-    "ip_address",
-    "postal_code",
-    "flight_number",
-    "license_plate_number",
-    "vehicle_registration_number",
-    "digital_signature"
+    "person", "organization", "social_media_handle", "username", "insurance_company",
+    "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number",
+    "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc",
+    "bank_account_number", "iban", "transaction_number", "cpf", "cnpj",
+    "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number", "identity_card_number",
+    "national_id_number", "identity_document_number", "birth_certificate_number", "social_security_number",
+    "health_insurance_id_number", "health_insurance_number", "national_health_insurance_number",
+    "student_id_number", "registration_number", "insurance_number", "serial_number",
+    "visa_number", "reservation_number", "train_ticket_number",
+    "medication", "medical_condition", "blood_type", "date_of_birth",
+    "address", "ip_address", "postal_code",
+    "flight_number", "license_plate_number", "vehicle_registration_number", "digital_signature"
 ]
+
 # Corrected mapping dictionary
 category_mapping = {
     "People_and_Groups": [
-        "person",
-        "organization",
-        "social_media_handle",
-        "username",
-        "insurance_company"
+        "person", "organization", "social_media_handle", "username", "insurance_company"
     ],
     "Contact_Information": [
-        "phone_number",
-        "email",
-        "email_address",
-        "mobile_phone_number",
-        "landline_phone_number",
-        "fax_number"
+        "phone_number", "email", "email_address", "mobile_phone_number", "landline_phone_number", "fax_number"
     ],
     "Financial_and_Transactions": [
-        "credit_card_number",
-        "credit_card_expiration_date",
-        "credit_card_brand",
-        "cvv",
-        "cvc",
-        "bank_account_number",
-        "iban",
-        "transaction_number",
-        "cpf",
-        "cnpj"
+        "credit_card_number", "credit_card_expiration_date", "credit_card_brand", "cvv", "cvc",
+        "bank_account_number", "iban", "transaction_number", "cpf", "cnpj"
     ],
     "Identification_and_Documents": [
-        "passport_number",
-        "passport_expiration_date",
-        "driver's_license_number",
-        "tax_identification_number",
-        "identity_card_number",
-        "national_id_number",
-        "identity_document_number",
-        "birth_certificate_number",
-        "social_security_number",
-        "health_insurance_id_number",
-        "health_insurance_number",
-        "national_health_insurance_number",
-        "student_id_number",
-        "registration_number",
-        "insurance_number",
-        "serial_number",
-        "visa_number",
-        "reservation_number",
-        "train_ticket_number"
+        "passport_number", "passport_expiration_date", "driver's_license_number", "tax_identification_number",
+        "identity_card_number", "national_id_number", "identity_document_number", "birth_certificate_number",
+        "social_security_number", "health_insurance_id_number", "health_insurance_number",
+        "national_health_insurance_number", "student_id_number", "registration_number",
+        "insurance_number", "serial_number", "visa_number", "reservation_number", "train_ticket_number"
     ],
     "Health_and_Personal": [
-        "medication",
-        "medical_condition",
-        "blood_type",
-        "date_of_birth"
+        "medication", "medical_condition", "blood_type", "date_of_birth"
     ],
     "Locations_and_Addresses": [
-        "address",
-        "ip_address",
-        "postal_code"
+        "address", "ip_address", "postal_code"
     ],
     "Transportation_and_Logistics": [
-        "flight_number",
-        "license_plate_number",
-        "vehicle_registration_number"
+        "flight_number", "license_plate_number", "vehicle_registration_number"
     ],
     "Digital_and_Security": [
         "digital_signature"
@@ -273,7 +204,6 @@ if st.button("Results"):
         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
@@ -282,6 +212,8 @@ if st.button("Results"):
                 entities = model.predict_entities(text, labels)
                 df = pd.DataFrame(entities)
                 st.session_state.results_df = df
+                
+                # Move the Comet ML logging and termination here
                 if not df.empty:
                     df['category'] = df['label'].map(reverse_category_mapping)
                     if comet_initialized:
@@ -293,15 +225,16 @@ if st.button("Results"):
                         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
         else:
-            # If the text is the same, just show the cached results without re-running
             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")
@@ -395,7 +328,7 @@ if st.session_state.show_results:
         
         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: #81C784; border: 1px solid black; padding: 5px; color: #1B5E20; }""",
         ):
             st.download_button(
                 label="Download results and glossary (zip)",
@@ -403,14 +336,10 @@ if st.session_state.show_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()
-    elapsed_time = end_time - start_time
+        
     st.text("")
     st.text("")
-    st.info(f"Results processed in **{elapsed_time:.2f} seconds**.") 
+    st.info(f"Results processed in **{st.session_state.elapsed_time:.2f} seconds**.")