Update src/streamlit_app.py

src/streamlit_app.py

@@ -8,8 +8,6 @@ import plotly.express as px
 import zipfile
 import json
 import hashlib
-
-
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
@@ -18,8 +16,6 @@ from comet_ml import Experiment
 st.set_page_config(layout="wide", page_title="NER")
 st.subheader("HR.ai", divider="green")
 st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
-
-
 st.markdown(
     """
     <style>
@@ -73,37 +69,9 @@ st.markdown(
     }
     </style>
     """,
-    unsafe_allow_html=True
-)
-
+    unsafe_allow_html=True)
 expander = st.expander("**Important notes**")
-expander.write(""" 
-**How to Use the HR.ai web app:** 
-1. Type or paste your text into the text area, then press Ctrl + Enter.
-2. Click the 'Results' button to extract and tag entities in your text data. 
-
-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 the Question-Answering feature:**
-1. Type or paste your text into the text area, then press Ctrl + Enter. 
-2. Click the 'Add Question' button to add your question to the Record of Questions. You can manage your questions by deleting them one by one.
-3. Click the 'Extract Answers' button to extract the answer to your question.
-
-Results are presented in an easy-to-read table, visualized in an interactive tree map, and is available for download.   
-
-
-**Entities:** "Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Job_title", "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"
-
-**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(""" **How to Use the HR.ai web app:** 1. Type or paste your text into the text area, then press Ctrl + Enter.2. Click the 'Results' button to extract and tag entities in your text data. 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 the Question-Answering feature:**1. Type or paste your text into the text area, then press Ctrl + Enter. 2. Click the 'Add Question' button to add your question to the Record of Questions. You can manage your questions by deleting them one by one.3. Click the 'Extract Answers' button to extract the answer to your question.Results are presented in an easy-to-read table, visualized in an interactive tree map, and is available for download.   **Entities:** "Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Job_title", "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"**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 web app on your website. Feel free to adjust the width and height values to fit your page.")
     code = '''
@@ -126,12 +94,11 @@ 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.")
 
 # --- Model Loading and Caching ---
-@st.cache_resource
+@st.cache_resourced
 def load_gliner_model(model_name):
     """Initializes and caches the GLiNER model."""
     try:
@@ -145,7 +112,6 @@ def load_gliner_model(model_name):
 
 # --- HR_AI Model Labels and Mappings ---
 labels = ["Email", "Phone_number", "Street_address", "City", "Country", "Date_of_birth", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Terminated", "Retired", "Job_title", "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"]
-
 category_mapping = {
     "Contact Information": ["Email", "Phone_number", "Street_address", "City", "Country"],
     "Personal Details": ["Date_of_birth", "Marital_status", "Person"],
@@ -158,8 +124,7 @@ category_mapping = {
     "Deductions": ["Tax", "Deductions"],
     "Recruitment & Sourcing": ["Interview_type", "Applicant", "Referral", "Job_board", "Recruiter"],
     "Legal & Compliance": ["Offer_letter", "Agreement"],
-    "Professional_Development": ["Certification", "Skill"]
-}
+    "Professional_Development": ["Certification", "Skill"]}
 reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list}
 
 # --- InfoFinder Helpers ---
@@ -173,17 +138,17 @@ def get_stable_color(label):
 
 # --- Main App with Tabs ---
 tab1, tab2 = st.tabs(["HR.ai", "Question-Answering"])
-
 with tab1:
-    
-    
     # Load model for this tab
     model_hr = load_gliner_model("HR_AI")
-
     text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area_hr')
 
     def clear_text_hr():
         st.session_state['my_text_area_hr'] = ""
+        st.session_state.hr_data_to_download = None
+
+    if 'hr_data_to_download' not in st.session_state:
+        st.session_state.hr_data_to_download = None
 
     st.button("Clear text", on_click=clear_text_hr, key="clear_hr")
 
@@ -191,6 +156,7 @@ with tab1:
         start_time = time.time()
         if not text.strip():
             st.warning("Please enter some text to extract entities.")
+            st.session_state.hr_data_to_download = None
         else:
             with st.spinner("Extracting entities...", show_time=True):
                 entities = model_hr.predict_entities(text, labels)
@@ -203,103 +169,108 @@ with tab1:
                         experiment.log_parameter("input_text", text)
                         experiment.log_table("predicted_entities", df)
 
-                    st.subheader("Grouped Entities by Category", divider="green")
-                    category_names = sorted(list(category_mapping.keys()))
-                    category_tabs_hr = st.tabs(category_names)
-                    for i, category_name in enumerate(category_names):
-                        with category_tabs_hr[i]:
-                            df_category_filtered = df[df['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']
-                        - **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(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='#F5FFFA', plot_bgcolor='#F5FFFA')
-                    st.plotly_chart(fig_treemap)
-                    
-                    col1, col2 = st.columns(2)
-                    with col1:
-                        st.subheader("Pie chart", divider="green")
-                        grouped_counts = df['category'].value_counts().reset_index()
-                        grouped_counts.columns = ['category', 'count']
-                        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(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
-                        st.plotly_chart(fig_pie)
-                    
-                    with col2:
-                        st.subheader("Bar chart", divider="green")
-                        fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
-                        fig_bar.update_layout(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
-                        st.plotly_chart(fig_bar)
-                    
-                    st.subheader("Most Frequent Entities", divider="green")
-                    word_counts = df['text'].value_counts().reset_index()
-                    word_counts.columns = ['Entity', 'Count']
-                    repeating_entities = word_counts[word_counts['Count'] > 1]
-                    if not repeating_entities.empty:
-                        st.dataframe(repeating_entities, use_container_width=True)
-                        fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
-                        fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'}, paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
-                        st.plotly_chart(fig_repeating_bar)
-                    else:
-                        st.warning("No entities were found that occur more than once.")
-                    
-                    st.divider()
-                    
+                    # Prepare data for download and store it in session state
                     dfa = pd.DataFrame(data={'Column Name': ['text', 'label', 'score', 'start', 'end'], 'Description': ['entity extracted from your text data', 'label (tag) assigned to a given extracted entity', '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']})
                     buf = io.BytesIO()
                     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))
+                    st.session_state.hr_data_to_download = buf.getvalue()
                     
-                    st.download_button(
-                        label="Download results and glossary (zip)",
-                        data=buf.getvalue(),
-                        file_name="nlpblogs_results.zip",
-                        mime="application/zip",
-                    )
-                    
-                    if comet_initialized:
-                        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
-                        experiment.end()
+                    st.session_state.hr_df = df
                 else:
+                    st.session_state.hr_data_to_download = None
                     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**.")
 
-with tab2:
-    
+                end_time = time.time()
+                elapsed_time = end_time - start_time
+                st.text("")
+                st.text("")
+                st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
+
+    # Display logic for HR tab (always runs if data is in session state)
+    if 'hr_df' in st.session_state and not st.session_state.hr_df.empty:
+        df = st.session_state.hr_df
+        st.subheader("Grouped Entities by Category", divider="green")
+        category_names = sorted(list(category_mapping.keys()))
+        category_tabs_hr = st.tabs(category_names)
+        for i, category_name in enumerate(category_names):
+            with category_tabs_hr[i]:
+                df_category_filtered = df[df['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']
+            - **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(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='#F5FFFA', plot_bgcolor='#F5FFFA')
+        st.plotly_chart(fig_treemap)
+
+        col1, col2 = st.columns(2)
+        with col1:
+            st.subheader("Pie chart", divider="green")
+            grouped_counts = df['category'].value_counts().reset_index()
+            grouped_counts.columns = ['category', 'count']
+            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(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
+            st.plotly_chart(fig_pie)
+
+        with col2:
+            st.subheader("Bar chart", divider="green")
+            fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
+            fig_bar.update_layout(paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
+            st.plotly_chart(fig_bar)
+
+        st.subheader("Most Frequent Entities", divider="green")
+        word_counts = df['text'].value_counts().reset_index()
+        word_counts.columns = ['Entity', 'Count']
+        repeating_entities = word_counts[word_counts['Count'] > 1]
+        if not repeating_entities.empty:
+            st.dataframe(repeating_entities, use_container_width=True)
+            fig_repeating_bar = px.bar(repeating_entities, x='Entity', y='Count', color='Entity')
+            fig_repeating_bar.update_layout(xaxis={'categoryorder': 'total descending'}, paper_bgcolor='#F5FFFA', plot_bgcolor='#F5FFFA')
+            st.plotly_chart(fig_repeating_bar)
+        else:
+            st.warning("No entities were found that occur more than once.")
+
+        st.divider()
+        if st.session_state.hr_data_to_download:
+            st.download_button(
+                label="Download results and glossary (zip)",
+                data=st.session_state.hr_data_to_download,
+                file_name="nlpblogs_results.zip",
+                mime="application/zip",
+            )
+            if comet_initialized:
+                experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
+                experiment.end()
 
+with tab2:
     # Load model for this tab
     model_qa = load_gliner_model("InfoFinder")
-
     user_text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area_infofinder')
-
+    
     def clear_text_qa():
         st.session_state['my_text_area_infofinder'] = ""
+        st.session_state.user_labels = []
+        st.session_state.qa_data_to_download = None
 
     st.button("Clear text", on_click=clear_text_qa, key="clear_qa")
-
     st.subheader("Question-Answering", divider="green")
     question_input = st.text_input("Ask wh-questions. **Wh-questions begin with what, when, where, who, whom, which, whose, why and how. We use them to ask for specific information.**")
-
+    
     if st.button("Add Question"):
         if question_input:
             if question_input not in st.session_state.user_labels:
@@ -309,10 +280,10 @@ with tab2:
                 st.warning("This question has already been added.")
         else:
             st.warning("Please enter a 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])
@@ -324,9 +295,9 @@ with tab2:
                     st.rerun()
     else:
         st.info("No questions defined yet. Use the input above to add one.")
-
+    
     st.divider()
-
+    
     if st.button("Extract Answers"):
         if not user_text.strip():
             st.warning("Please enter some text to analyze.")
@@ -337,7 +308,7 @@ with tab2:
                 experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=COMET_PROJECT_NAME)
                 experiment.log_parameter("input_text_length", len(user_text))
                 experiment.log_parameter("defined_labels", st.session_state.user_labels)
-
+            
             start_time = time.time()
             with st.spinner("Analyzing text...", show_time=True):
                 try:
@@ -345,46 +316,45 @@ with tab2:
                     end_time = time.time()
                     elapsed_time = end_time - start_time
                     st.info(f"Processing took **{elapsed_time:.2f} seconds**.")
-
+                    
                     if entities:
                         df1 = pd.DataFrame(entities)
                         df2 = df1[['label', 'text', 'score']]
                         df = df2.rename(columns={'label': 'question', 'text': 'answer'})
+                        st.session_state.qa_data_to_download = df.to_csv(index=False).encode('utf-8')
+                        st.session_state.qa_df = df
                         
-                        st.subheader("Extracted Answers", divider="green")
-                        st.dataframe(df, use_container_width=True)
-
-                        st.subheader("Tree map", divider="green")
-                        all_labels = df['question'].unique()
-                        label_color_map = {label: get_stable_color(label) for label in all_labels}
-                        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'question', 'answer'], values='score', color='question', color_discrete_map=label_color_map)
-                        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#F3E5F5', plot_bgcolor='#F3E5F5')
-                        st.plotly_chart(fig_treemap)
-
-                        csv_data = df.to_csv(index=False).encode('utf-8')
-                        st.download_button(
-                            label="Download CSV",
-                            data=csv_data,
-                            file_name="nlpblogs_questions_answers.csv",
-                            mime="text/csv",
-                        )
-
                         if comet_initialized:
                             experiment.log_metric("processing_time_seconds", elapsed_time)
                             experiment.log_table("predicted_entities", df)
-                            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")
-                            experiment.end()
                     else:
-                        st.info("No answers were found in the text with the defined questions.")
-                        if comet_initialized:
-                            experiment.end()
+                        st.session_state.qa_data_to_download = None
+                        st.warning("No answers were found in the text with the defined questions.")
                 except Exception as e:
                     st.error(f"An error occurred during processing: {e}")
                     st.write(f"Error details: {e}")
+                finally:
                     if comet_initialized:
-                        experiment.log_text(f"Error: {e}")
                         experiment.end()
-
-      
-
-            
+    
+    # Display logic for QA tab (always runs if data is in session state)
+    if 'qa_df' in st.session_state and not st.session_state.qa_df.empty:
+        df = st.session_state.qa_df
+        st.subheader("Extracted Answers", divider="green")
+        st.dataframe(df, use_container_width=True)
+        st.subheader("Tree map", divider="green")
+        all_labels = df['question'].unique()
+        label_color_map = {label: get_stable_color(label) for label in all_labels}
+        fig_treemap = px.treemap(df, path=[px.Constant("all"), 'question', 'answer'], values='score', color='question', color_discrete_map=label_color_map)
+        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25), paper_bgcolor='#F3E5F5', plot_bgcolor='#F3E5F5')
+        st.plotly_chart(fig_treemap)
+        
+        if st.session_state.qa_data_to_download:
+            st.download_button(
+                label="Download CSV",
+                data=st.session_state.qa_data_to_download,
+                file_name="nlpblogs_questions_answers.csv",
+                mime="text/csv",
+            )
+            if comet_initialized:
+                experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")