Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,6 +1,3 @@
-import os
-os.environ['HF_HOME'] = '/tmp'
-
 import os
 import time
 import streamlit as st
@@ -18,6 +15,9 @@ import hashlib
 # Set up environment variables
 os.environ['HF_HOME'] = '/tmp'
 
+# --- Page Configuration and UI Elements ---
+st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
+
 st.markdown(
     """
     <style>
@@ -69,10 +69,9 @@ st.markdown(
     unsafe_allow_html=True
 )
 
-# --- Page Configuration and UI Elements ---
-st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
 st.subheader("HR.ai", divider="green")
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
+
 expander = st.expander("**Important notes**")
 expander.write("""**Named Entities:** This HR.ai predicts thirty-five (35) labels: "Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Date", "Organization", "Role", "Performance_score", "Leave_of_absence", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance", "Pay_rate", "Annual_salary", "Tax", "Deductions", "Interview_type", "Applicant", "Referral", "Job_board", "Recruiter", "Offer_letter", "Agreement", "Certification", "Skill"
 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.
@@ -98,7 +97,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.")
 
@@ -122,7 +120,7 @@ category_mapping = {
 }
 
 # --- Model Loading ---
-@st.cache_resource
+@st.cache_resourced
 def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
@@ -170,34 +168,34 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df_ner)
-                
-                st.subheader("Grouped Entities by Category", divider="green")
-                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_ner[df_ner['category'] == category_name]
-                        if not df_category_filtered.empty:
-                            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']
-                    - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
-                    - **label**: ['label (tag) assigned to a given extracted entity']
-                    - **category**: ['the high-level category for the label']
-                    - **start**: ['index of the start of the corresponding entity']
-                    - **end**: ['index of the end of the corresponding entity']
-                    ''')
             else:
                 st.warning("No entities were found in the provided text.")
                 if 'df_ner' in st.session_state:
                     del st.session_state.df_ner
 
-# --- Treemap Display Section ---
+# --- Display Sections based on Session State ---
 if 'df_ner' in st.session_state and not st.session_state.df_ner.empty:
+    st.subheader("Grouped Entities by Category", divider="green")
+    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 = st.session_state.df_ner[st.session_state.df_ner['category'] == category_name]
+            if not df_category_filtered.empty:
+                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']
+        - **score**: ['accuracy score; how accurately a tag has been assigned to a given entity']
+        - **label**: ['label (tag) assigned to a given extracted entity']
+        - **category**: ['the high-level category for the label']
+        - **start**: ['index of the start of the corresponding entity']
+        - **end**: ['index of the end of the corresponding entity']
+        ''')
+
     st.divider()
     st.subheader("Candidate Card", divider="green")
     fig_treemap = px.treemap(st.session_state.df_ner, path=[px.Constant("all"), 'category', 'label', 'text'], values='score', color='category')
@@ -215,6 +213,7 @@ def load_gliner_model():
         st.stop()
 
 qa_model = load_gliner_model()
+
 st.subheader("Question-Answering", divider="green")
 
 if 'user_labels' not in st.session_state:
@@ -235,6 +234,7 @@ if st.button("Add Question"):
 
 st.markdown("---")
 st.subheader("Record of Questions", divider="green")
+
 if st.session_state.user_labels:
     for i, label in enumerate(st.session_state.user_labels):
         col_list, col_delete = st.columns([0.9, 0.1])
@@ -271,19 +271,16 @@ if st.button("Extract Answers"):
                 end_time = time.time()
                 elapsed_time = end_time - start_time
                 st.info(f"Processing took **{elapsed_time:.2f} seconds**.")
-
                 if entities:
                     df_qa = pd.DataFrame(entities)
                     df_qa = df_qa[['label', 'text', 'score']].rename(columns={'label': 'question', 'text': 'answer'})
                     st.session_state.df_qa = df_qa  # Store QA results in session state
-
                     st.subheader("Extracted Answers", divider="green")
                     st.dataframe(df_qa, use_container_width=True)
                 else:
                     st.warning("No answers were found for the provided questions.")
                     if 'df_qa' in st.session_state:
                         del st.session_state.df_qa
-
             except Exception as e:
                 st.error(f"An error occurred during answer extraction: {e}")
                 if 'df_qa' in st.session_state:
@@ -292,7 +289,6 @@ if st.button("Extract Answers"):
 # --- Download Button Section ---
 def create_zip_file_and_get_bytes():
     """Generates a zip file in memory with all available dataframes."""
-    
     # Define the glossary DataFrame here to ensure it's always available
     dfa = pd.DataFrame(
         data={
@@ -307,10 +303,8 @@ def create_zip_file_and_get_bytes():
             ]
         }
     )
-
     if 'df_ner' not in st.session_state and 'df_qa' not in st.session_state:
         return None, None
-
     buf = io.BytesIO()
     with zipfile.ZipFile(buf, "w") as myzip:
         if 'df_ner' in st.session_state and not st.session_state.df_ner.empty:
@@ -318,7 +312,6 @@ def create_zip_file_and_get_bytes():
         if 'df_qa' in st.session_state and not st.session_state.df_qa.empty:
             myzip.writestr("Extracted_Answers.csv", st.session_state.df_qa.to_csv(index=False))
         myzip.writestr("Glossary_of_tags.csv", dfa.to_csv(index=False))
-
     return buf.getvalue(), "nlpblogs_results.zip"
 
 st.divider()
@@ -336,5 +329,4 @@ if ('df_ner' in st.session_state and not st.session_state.df_ner.empty) or \
                 data=zip_data,
                 file_name=file_name,
                 mime="application/zip",
-            )
-
+            )