Update src/streamlit_app.py

src/streamlit_app.py

@@ -7,96 +7,119 @@ import io
 import plotly.express as px
 import zipfile
 import json
+import hashlib
+
 from streamlit_extras.stylable_container import stylable_container
 from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
-import hashlib
 
 # --- Page Configuration and UI Elements ---
-st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
+st.set_page_config(layout="wide", page_title="Combined NLP App")
+st.subheader("Combined NLP App", divider="green")
+st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
+
 st.markdown(
     """
     <style>
-    /* Main app background and text color */
     .stApp {
-        background-color: #F5FFFA; /* Mint cream, a very light green */
-        color: #000000; /* Black for the text */
+        background-color: #F5FFFA;
+        color: #000000;
     }
-    /* Sidebar background color */
     .css-1d36184 {
-        background-color: #B2F2B2; /* A pale green for the sidebar */
+        background-color: #B2F2B2;
         secondary-background-color: #B2F2B2;
     }
-    /* Expander background color */
     .streamlit-expanderContent {
         background-color: #F5FFFA;
     }
-    /* Expander header background color */
     .streamlit-expanderHeader {
         background-color: #F5FFFA;
     }
-    /* Text Area background and text color */
     .stTextArea textarea {
-        background-color: #D4F4D4; /* A light, soft green */
-        color: #000000; /* Black for text */
+        background-color: #D4F4D4;
+        color: #000000;
     }
-    /* Text input background and text color */
-    .stTextInput textinput {
-        background-color: #D4F4D4; /* A light, soft green */
-        color: #000000; /* Black for text */
+    .stTextInput input {
+        background-color: #B39DDB;
+        color: #1A0A26;
     }
-    /* Button background and text color */
     .stButton > button {
         background-color: #D4F4D4;
         color: #000000;
     }
-    /* Warning box background and text color */
     .stAlert.st-warning {
-        background-color: #C8F0C8; /* A light green for the warning box */
+        background-color: #C8F0C8;
         color: #000000;
     }
-    /* Success box background and text color */
     .stAlert.st-success {
-        background-color: #C8F0C8; /* A light green for the success box */
+        background-color: #C8F0C8;
         color: #000000;
     }
+    .stTabs [data-baseweb="tab-list"] button [data-testid="stMarkdownContainer"] p {
+        font-size: 1.25rem;
+    }
     </style>
     """,
-    unsafe_allow_html=True)
-st.subheader("HR.ai", divider="green")
-st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
+    unsafe_allow_html=True
+)
+
 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.**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("""
+This application combines two NLP tools in one place:
+- **HR.ai**: Specializes in extracting 36 predefined HR-related entities from your text.
+- **InfoFinder**: Answers custom wh-questions by finding relevant entities in your text.
+
+Both tools feature interactive visualizations and data download options.
+""")
+
 with st.sidebar:
-    st.write("Use the following code to embed the HR.ai web app on your website. Feel free to adjust the width and height values to fit your page.")
+    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 = '''
-    <iframe	src="https://aiecosystem-hr-ai.hf.space"	frameborder="0"	width="850"	height="450"    ></iframe>
+    <iframe
+    src="https://aiecosystem-hr-ai.hf.space"
+    frameborder="0"
+    width="850"
+    height="450"
+    ></iframe>
     '''
     st.code(code, language="html")
     st.text("")
     st.text("")
     st.divider()
     st.subheader("🚀 Ready to build your own AI Web App?", divider="green")
-    st.link_button("AI Web App Builder", " https://nlpblogs.com/custom-web-app-development/", type="primary")
+    st.link_button("AI Web App Builder", "https://nlpblogs.com/custom-web-app-development/", 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 = ["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"]
+# --- Model Loading and Caching ---
+@st.cache_resource
+def load_gliner_model(model_name):
+    """Initializes and caches the GLiNER model."""
+    try:
+        if model_name == "HR_AI":
+            return GLiNER.from_pretrained("knowledgator/gliner-multitask-large-v0.5", nested_ner=True, num_gen_sequences=2, gen_constraints=labels)
+        elif model_name == "InfoFinder":
+            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()
+
+# --- 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"]
 
-# Create a mapping dictionary for labels to categories
 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": ["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"],
@@ -104,218 +127,234 @@ 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}
 
-# --- Model Loading ---
-@st.cache_resource
-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)
-    except Exception as e:
-        st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
-        st.stop()
-model = load_ner_model()
+# --- InfoFinder Helpers ---
+if 'user_labels' not in st.session_state:
+    st.session_state.user_labels = []
 
-# 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}
+def get_stable_color(label):
+    hash_object = hashlib.sha1(label.encode('utf-8'))
+    hex_dig = hash_object.hexdigest()
+    return '#' + hex_dig[:6]
 
-# --- 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')
-
-def clear_text():
-    """Clears the text area and session state."""
-    st.session_state['my_text_area'] = ""
-    # Clear stored results
-    for key in ['df_ner', 'df_qa', 'fig_treemap', 'user_labels']:
-        if key in st.session_state:
-            del st.session_state[key]
-    st.rerun()
-
-st.button("Clear text", on_click=clear_text)
-
-# --- Results Section for NER ---
-if st.button("Results"):
-    start_time = time.time()
-    if not text.strip():
-        st.warning("Please enter some text to extract entities.")
-    else:
-        with st.spinner("Extracting entities...", show_time=True):
-            entities = model.predict_entities(text, labels)
-            df_ner = pd.DataFrame(entities)
-            if not df_ner.empty:
-                df_ner['category'] = df_ner['label'].map(reverse_category_mapping)
-                st.session_state.df_ner = df_ner  # Store df in session state
-                if comet_initialized:
-                    experiment = Experiment(
-                        api_key=COMET_API_KEY,
-                        workspace=COMET_WORKSPACE,
-                        project_name=COMET_PROJECT_NAME,
-                    )
-                    experiment.log_parameter("input_text", text)
-                    experiment.log_table("predicted_entities", df_ner)
-            else:
-                st.warning("No entities were found in the provided text.")
-                if 'df_ner' in st.session_state:
-                    del st.session_state.df_ner
-
-# --- 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')
-    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)
-    df_ner_results = st.session_state.df_ner.drop(columns=['category']) # Define df_ner_results here
-    dfa = pd.DataFrame(
-        data={
-            'Column Name': ['text', 'label', 'score', 'start', 'end', 'category'],
-            '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',
-                            'the broader category the entity belongs to',]}
-    )
-    buf = io.BytesIO()
-    with zipfile.ZipFile(buf, "w") as myzip:
-        myzip.writestr("Summary of the results.csv", df_ner_results.to_csv(index=False)) # Use df_ner_results
-        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; }""",
-    ):
-        st.download_button(
-        label="Download results and glossary (zip)",
-        data=buf.getvalue(),
-        file_name="nlpblogs_results.zip",
-        mime="application/zip",)
+# --- Main App with Tabs ---
+tab1, tab2 = st.tabs(["HR.ai (Named Entity Recognition)", "InfoFinder (Question-Answering)"])
+
+with tab1:
+    st.markdown("### HR.ai: Named Entity Recognition")
     
-    if comet_initialized:
-        experiment = Experiment(
-            api_key=COMET_API_KEY,
-            workspace=COMET_WORKSPACE,
-            project_name=COMET_PROJECT_NAME,
-        )
-        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
-        experiment.end()
-    else:
-        st.warning("No entities were found in the provided text.")
+    # Load model for this tab
+    model_hr = load_gliner_model("HR_AI")
 
-# --- Question Answering Section ---
-@st.cache_resource
-def load_gliner_model():
-    """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()
+    text = st.text_area("Type or paste your text below, and then press Ctrl + Enter", height=250, key='my_text_area_hr')
 
-qa_model = load_gliner_model()
-st.subheader("Question-Answering", divider="green")
-if 'user_labels' not in st.session_state:
-    st.session_state.user_labels = []
+    def clear_text_hr():
+        st.session_state['my_text_area_hr'] = ""
 
-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.**")
+    st.button("Clear text", on_click=clear_text_hr, key="clear_hr")
 
-if st.button("Add Question"):
-    if question_input:
-        if question_input not in st.session_state.user_labels:
-            st.session_state.user_labels.append(question_input)
-            st.success(f"Added question: {question_input}")
+    if st.button("Extract HR Entities"):
+        start_time = time.time()
+        if not text.strip():
+            st.warning("Please enter some text to extract entities.")
         else:
-            st.warning("This question has already been added.")
-    else:
-        st.warning("Please enter a question.")
-    st.rerun()
-
-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])
-        with col_list:
-            st.write(f"- {label}", key=f"label_{i}")
-        with col_delete:
-            if st.button("Delete", key=f"delete_{i}"):
-                st.session_state.user_labels.pop(i)
-                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 text.strip():
-        st.warning("Please enter some text to analyze.")
-    elif not st.session_state.user_labels:
-        st.warning("Please define at least one question.")
+            with st.spinner("Extracting entities...", show_time=True):
+                entities = model_hr.predict_entities(text, labels)
+                df = pd.DataFrame(entities)
+
+                if not df.empty:
+                    df['category'] = df['label'].map(reverse_category_mapping)
+                    if comet_initialized:
+                        experiment = Experiment(api_key=COMET_API_KEY, workspace=COMET_WORKSPACE, project_name=COMET_PROJECT_NAME)
+                        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']
+                        - **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("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)
+                    
+                    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()
+                    
+                    dfa = pd.DataFrame(data={'Column Name': ['text', 'label', 'score', 'start', 'end', 'category'], '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', 'the broader category the entity belongs to']})
+                    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.download_button(
+                        label="Download results and glossary (zip)",
+                        data=buf.getvalue(),
+                        file_name="nlpblogs_hr_results.zip",
+                        mime="application/zip",
+                    )
+                    
+                    if comet_initialized:
+                        experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap_categories")
+                        experiment.end()
+                else:
+                    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:
+    st.markdown("### InfoFinder: Question-Answering")
+
+    # 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.button("Clear text", on_click=clear_text_qa, key="clear_qa")
+
+    st.subheader("Question-Answering", divider="violet")
+    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:
+                st.session_state.user_labels.append(question_input)
+                st.success(f"Added question: {question_input}")
+            else:
+                st.warning("This question has already been added.")
+        else:
+            st.warning("Please enter a question.")
+            
+    st.markdown("---")
+    st.subheader("Record of Questions", divider="violet")
+
+    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])
+            with col_list:
+                st.write(f"- {label}", key=f"label_{i}")
+            with col_delete:
+                if st.button("Delete", key=f"delete_{i}"):
+                    st.session_state.user_labels.pop(i)
+                    st.rerun()
     else:
-        if comet_initialized:
-            experiment = Experiment(
-                api_key=COMET_API_KEY,
-                workspace=COMET_WORKSPACE,
-                project_name=COMET_PROJECT_NAME
-            )
-            experiment.log_parameter("input_text_length", len(text))
-            experiment.log_parameter("defined_labels", st.session_state.user_labels)
+        st.info("No questions defined yet. Use the input above to add one.")
 
-        start_time = time.time()
-        with st.spinner("Analyzing text...", show_time=True):
-            try:
-                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**.")
-                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)
-                    csv_data = df_qa.to_csv(index=False).encode('utf-8')
-                    with stylable_container(
-                        key="download_button_qa",
-                        css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
-                    ):
+    st.divider()
+
+    if st.button("Extract Answers"):
+        if not user_text.strip():
+            st.warning("Please enter some text to analyze.")
+        elif not st.session_state.user_labels:
+            st.warning("Please define at least one question.")
+        else:
+            if comet_initialized:
+                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:
+                    entities = model_qa.predict_entities(user_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**.")
+
+                    if entities:
+                        df1 = pd.DataFrame(entities)
+                        df2 = df1[['label', 'text', 'score']]
+                        df = df2.rename(columns={'label': 'question', 'text': 'answer'})
+                        
+                        st.subheader("Extracted Answers", divider="violet")
+                        st.dataframe(df, use_container_width=True)
+
+                        st.subheader("Tree map", divider="violet")
+                        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_extracted_answers.csv",
+                            file_name="nlpblogs_infofinder_results.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()
+                except Exception as e:
+                    st.error(f"An error occurred during processing: {e}")
+                    st.write(f"Error details: {e}")
                     if comet_initialized:
-                        experiment.log_metric("processing_time_seconds", elapsed_time)
-                        experiment.log_table("predicted_entities", df_qa)
-                        experiment.end()
-                else:
-                    st.info("No answers were found in the text with the defined questions.")
-                    if comet_initialized:
+                        experiment.log_text(f"Error: {e}")
                         experiment.end()
-            except Exception as e:
-                st.error(f"An error occurred during processing: {e}")
-                st.write(f"Error details: {e}")
-                if comet_initialized:
-                    experiment.log_text(f"Error: {e}")
-                    experiment.end()
+
+      
+
+