Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,7 +1,7 @@
 import os
 os.environ['HF_HOME'] = '/tmp'
 
-
+import os
 import time
 import streamlit as st
 import pandas as pd
@@ -16,6 +16,9 @@ from gliner import GLiNER
 from comet_ml import Experiment
 import hashlib
 
+# Set up environment variables
+os.environ['HF_HOME'] = '/tmp'
+
 st.markdown(
     """
     <style>
@@ -87,7 +90,6 @@ with st.sidebar:
     st.link_button("AI Web App Builder", " https://nlpblogs.com/custom-web-app-development/", type="primary")
 
 # --- Comet ML Setup ---
-os.environ['HF_HOME'] = '/tmp'
 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")
@@ -104,7 +106,7 @@ category_mapping = {
     "Contact Information": ["Email", "Phone_number", "Street_address", "City", "Country"],
     "Personal Details": ["Date_of_birth", "Marital_status", "Person"],
     "Employment Status": ["Full_time", "Part_time", "Contract", "Terminated", "Retired"],
-    "Employment Information" : ["Job_title", "Date", "Organization", "Role"],
+    "Employment Information": ["Job_title", "Date", "Organization", "Role"],
     "Performance": ["Performance_score"],
     "Attendance": ["Leave_of_absence"],
     "Benefits": ["Retirement_plan", "Bonus", "Stock_options", "Health_insurance"],
@@ -112,7 +114,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"]
 }
 
 # --- Model Loading ---
@@ -120,7 +122,7 @@ category_mapping = {
 def load_ner_model():
     """Loads the GLiNER model and caches it."""
     try:
-        return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints= labels)
+        return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
     except Exception as e:
         st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
         st.stop()
@@ -134,8 +136,13 @@ reverse_category_mapping = {label: category for category, label_list in category
 text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area')
 
 def clear_text():
-    """Clears the text area."""
+    """Clears the text area and session state."""
     st.session_state['my_text_area'] = ""
+    # Clear stored results
+    if 'df' in st.session_state:
+        del st.session_state.df
+    if 'fig_treemap' in st.session_state:
+        del st.session_state.fig_treemap
 
 st.button("Clear text", on_click=clear_text)
 
@@ -150,6 +157,7 @@ if st.button("Results"):
             df = pd.DataFrame(entities)
             if not df.empty:
                 df['category'] = df['label'].map(reverse_category_mapping)
+                st.session_state.df = df # Store df in session state
                 if comet_initialized:
                     experiment = Experiment(
                         api_key=COMET_API_KEY,
@@ -158,12 +166,10 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df)
-
+                
                 st.subheader("Grouped Entities by Category", divider="green")
-                # 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]
@@ -181,50 +187,39 @@ if st.button("Results"):
                     - **start**: ['index of the start of the corresponding entity']
                     - **end**: ['index of the end of the corresponding entity']
                     ''')
-                st.divider()
-
-                # Tree map
-                st.subheader("Tree map", 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)
-
-                if comet_initialized:
-                    experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
-                    experiment.end()
 
-            else:  # If df is empty from the first extraction
+            else:
                 st.warning("No entities were found in the provided text.")
+                # Clear session state if no results found
+                if 'df' in st.session_state:
+                    del st.session_state.df
 
-        end_time = time.time()
-        elapsed_time = end_time - start_time
-        st.text("")
-        st.text("")
-        st.info(f"Results processed in **{elapsed_time:.2f} seconds**.")
+# --- Treemap Display Section ---
+if 'df' in st.session_state and not st.session_state.df.empty:
+    st.divider()
+    st.subheader("Tree map", divider="green")
+    fig_treemap = px.treemap(st.session_state.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)
 
-# --- Question Answering Section (Moved outside the "Results" button) ---
-# --- Model Loading and Caching ---
+# --- Question Answering Section ---
 @st.cache_resource
 def load_gliner_model():
-    """
-    Initializes and caches the GLiNER model.
-    This ensures the model is only loaded once, improving performance.
-    """
+    """Initializes and caches the GLiNER model for QA."""
     try:
         return GLiNER.from_pretrained("knowledgator/gliner-multitask-v1.0", device="cpu")
     except Exception as e:
         st.error(f"Error loading the GLiNER model: {e}")
         st.stop()
 
-# Load the model
-model = load_gliner_model()
+qa_model = load_gliner_model()
 st.subheader("Question-Answering", divider="violet")
-# Replaced two columns with a single text input
-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 'user_labels' not in st.session_state:
     st.session_state.user_labels = []
 
+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:
@@ -238,23 +233,19 @@ if st.button("Add Question"):
 st.markdown("---")
 st.subheader("Record of Questions", divider="violet")
 if st.session_state.user_labels:
-    # Use enumerate to create a unique key for each item
     for i, label in enumerate(st.session_state.user_labels):
         col_list, col_delete = st.columns([0.9, 0.1])
         with col_list:
             st.write(f"- {label}", key=f"label_{i}")
         with col_delete:
-            # Create a unique key for each button using the index
             if st.button("Delete", key=f"delete_{i}"):
-                # Remove the label at the specific index
                 st.session_state.user_labels.pop(i)
-                # Rerun to update the UI
                 st.rerun()
 else:
     st.info("No questions defined yet. Use the input above to add one.")
 
 st.divider()
-# --- Main Processing Logic ---
+
 if st.button("Extract Answers"):
     if not text.strip():
         st.warning("Please enter some text to analyze.")
@@ -269,10 +260,11 @@ if st.button("Extract Answers"):
             )
             experiment.log_parameter("input_text_length", len(text))
             experiment.log_parameter("defined_labels", st.session_state.user_labels)
+
         start_time = time.time()
         with st.spinner("Analyzing text...", show_time=True):
             try:
-                entities = model.predict_entities(text, st.session_state.user_labels)
+                entities = qa_model.predict_entities(text, st.session_state.user_labels)
                 end_time = time.time()
                 elapsed_time = end_time - start_time
                 st.info(f"Processing took **{elapsed_time:.2f} seconds**.")
@@ -314,16 +306,7 @@ if st.button("Extract Answers"):
                             file_name="nlpblogs_results.zip",
                             mime="application/zip",
                         )
-
-                    if comet_initialized:
-                        # Assuming fig_treemap is still defined from the main NER run
-                        # If not, you might need to re-generate it or handle the case where it's not available.
-                        try:
-                            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
-                        except NameError:
-                            pass # Or handle this gracefully
-                        experiment.end()
-                else:  # If df is empty
+                else:
                     st.warning("No answers were found for the provided questions.")
             except Exception as e:
                 st.error(f"An error occurred during answer extraction: {e}")