Update src/streamlit_app.py

src/streamlit_app.py

@@ -13,7 +13,6 @@ from typing import Optional
 from gliner import GLiNER
 from comet_ml import Experiment
 
-
 st.markdown(
     """
     <style>
@@ -27,7 +26,7 @@ st.markdown(
         background-color: #B2F2B2; /* A pale green for the sidebar */
         secondary-background-color: #B2F2B2;
     }
-            
+    
     /* Expander background color */
     .streamlit-expanderContent {
         background-color: #F5FFFA;
@@ -61,99 +60,55 @@ 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="orange")
+st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
+expander = st.expander("**Important notes**")
+expander.write("""**Named Entities:** This HR.ai predicts sixty (60) labels:"Email", "Phone_number", "Street_address", "City", "State", "Zip_code", "Country", "Date_of_birth", "Gender", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Temporary", "Terminated", "Active", "Retired", "Job_title", "Employment_type", "Year", "Date", "Company", "Organization", "Role", "Position","Performance_review", "Performance_rating", "Performance_score", "Sick_days", "Vacation_days", "Leave_of_absence", "Holidays", "Pension", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance", "Retire_date", "Pay_rate", "Hourly_wage", "Annual_salary", "Overtime_pay", "Tax", "Social_security", "Deductions", "Job_posting", "Job_description", "Interview_type", "Applicant", "Candidate", "Referral", "Job_board", "Recruiter","Contract", "Offer_letter", "Agreement", "Training_course", "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, German, French, Italian, Spanish, Portuguese      
 
-# --- Page Configuration and UI Elements ---
-st.set_page_config(layout="wide", page_title="Named Entity Recognition App")
+**Technical issues:** If your connection times out, please refresh the page or reopen the app's URL. 
 
-st.subheader("HR.ai", divider="orange")
-st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
-
-expander = st.expander("**Important notes**")
-expander.write("""
-   **Named Entities:** This HR.ai predicts sixty (60) labels:
-"Email", "Phone_number", "Street_address", "City", "State", "Zip_code", "Country", "Date_of_birth", "Gender", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Temporary", "Terminated", "Active", "Retired", "Job_title", "Employment_type", "Year", "Date", "Company", "Organization", "Role", "Position",
-"Performance_review", "Performance_rating", "Performance_score", "Sick_days", "Vacation_days", "Leave_of_absence", "Holidays", "Pension", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance","Retire_date",
-"Pay_rate", "Hourly_wage", "Annual_salary", "Overtime_pay", "Tax", "Social_security", "Deductions", "Job_posting", "Job_description", "Interview_type", "Applicant", "Candidate", "Referral", "Job_board", "Recruiter",
-"Contract", "Offer_letter", "Agreement", "Training_course", "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, German, French, Italian, Spanish, Portuguese
-   
-   **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 [email protected]
-""")
+For any errors or inquiries, please contact us at info@nlpblogs.com""")
 
 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.")
     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 NER Web App? 🚀", divider="orange")
     st.link_button("NER Builder", "https://nlpblogs.com", 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", "State", "Zip_code", "Country",
-"Date_of_birth", "Gender", "Marital_status", "Person",
-"Full_time", "Part_time", "Contract", "Temporary", "Terminated", "Active", "Retired",
-"Job_title", "Employment_type", "Year", "Date", "Company", "Organization", "Role", "Position",
-"Performance_review", "Performance_rating", "Performance_score",
-"Sick_days", "Vacation_days", "Leave_of_absence", "Holidays",
-"Pension", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance","Retire_date",
-"Pay_rate", "Hourly_wage", "Annual_salary", "Overtime_pay",
-"Tax", "Social_security", "Deductions",
-"Job_posting", "Job_description", "Interview_type", "Applicant", "Candidate", "Referral", "Job_board", "Recruiter",
-"Contract", "Offer_letter", "Agreement",
-"Training_course", "Certification", "Skill"]
-
-
-
+labels = ["Email", "Phone_number", "Street_address", "City", "State", "Zip_code", "Country", "Date_of_birth", "Gender", "Marital_status", "Person", "Full_time", "Part_time", "Contract", "Temporary", "Terminated", "Active", "Retired", "Job_title", "Employment_type", "Year", "Date", "Company", "Organization", "Role", "Position", "Performance_review", "Performance_rating", "Performance_score", "Sick_days", "Vacation_days", "Leave_of_absence", "Holidays", "Pension", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance", "Retire_date", "Pay_rate", "Hourly_wage", "Annual_salary", "Overtime_pay", "Tax", "Social_security", "Deductions", "Job_posting", "Job_description", "Interview_type", "Applicant", "Candidate", "Referral", "Job_board", "Recruiter", "Contract", "Offer_letter", "Agreement", "Training_course", "Certification", "Skill"]
 
 # Create a mapping dictionary for labels to categories
 category_mapping = {
-
-    
     "Contact Information": ["Email", "Phone_number", "Street_address", "City", "State", "Zip_code", "Country"],
     "Personal Details": ["Date_of_birth", "Gender", "Marital_status", "Person"],
     "Employment Status": ["Full_time", "Part_time", "Contract", "Temporary", "Terminated", "Active", "Retired"],
     "Employment Information" : ["Job_title", "Employment_type", "Year", "Date", "Company", "Organization", "Role", "Position"],
-    
     "Performance": ["Performance_review", "Performance_rating", "Performance_score"],
     "Attendance": ["Sick_days", "Vacation_days", "Leave_of_absence", "Holidays"],
     "Benefits": ["Pension", "Retirement_plan", "Bonus", "Stock_options", "Health_insurance","Retire_date"],
@@ -164,22 +119,17 @@ category_mapping = {
     "Professional_Development": ["Training_course", "Certification", "Skill"]
 }
 
-
 # --- 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)
-
+        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()
 
-
-
 # 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}
 
@@ -205,7 +155,6 @@ if st.button("Results"):
 
             if not df.empty:
                 df['category'] = df['label'].map(reverse_category_mapping)
-
                 if comet_initialized:
                     experiment = Experiment(
                         api_key=COMET_API_KEY,
@@ -214,9 +163,22 @@ if st.button("Results"):
                     )
                     experiment.log_parameter("input_text", text)
                     experiment.log_table("predicted_entities", df)
-
+                
                 st.subheader("Extracted Entities", divider = "orange")
-                st.dataframe(df.style.set_properties(**{"border": "2px solid gray", "color": "blue", "font-size": "16px"}))
+                
+                # 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]
+                        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.")
+
+                st.divider()
 
                 with st.expander("See Glossary of tags"):
                     st.write('''
@@ -227,34 +189,30 @@ 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 = "orange")
                 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))
                 st.plotly_chart(fig_treemap)
-
+                
                 # Pie and Bar charts
                 grouped_counts = df['category'].value_counts().reset_index()
                 grouped_counts.columns = ['category', 'count']
-
                 col1, col2 = st.columns(2)
+                
                 with col1:
                     st.subheader("Pie chart", divider = "orange")
-                    fig_pie = px.pie(grouped_counts, values='count', names='category',
-                                    hover_data=['count'], labels={'count': 'count'}, title='Percentage of predicted categories')
+                    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')
                     st.plotly_chart(fig_pie)
-
+                
                 with col2:
                     st.subheader("Bar chart", divider = "orange")
-                    fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True,
-                                    title='Occurrences of predicted categories')
+                    fig_bar = px.bar(grouped_counts, x="count", y="category", color="category", text_auto=True, title='Occurrences of predicted categories')
                     st.plotly_chart(fig_bar)
-
+                
                 # Most Frequent Entities
                 st.subheader("Most Frequent Entities", divider="orange")
                 word_counts = df['text'].value_counts().reset_index()
@@ -267,15 +225,10 @@ if st.button("Results"):
                     st.plotly_chart(fig_repeating_bar)
                 else:
                     st.warning("No entities were found that occur more than once.")
-
-
-
-
-   
-
+                
                 # Download Section
                 st.divider()
-                
+                                
                 dfa = pd.DataFrame(
                     data={
                         'Column Name': ['text', 'label', 'score', 'start', 'end', 'category'],
@@ -289,12 +242,11 @@ if st.button("Results"):
                         ]
                     }
                 )
-
                 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))
-
+                
                 with stylable_container(
                     key="download_button",
                     css_styles="""button { background-color: red; border: 1px solid black; padding: 5px; color: white; }""",
@@ -305,17 +257,15 @@ if st.button("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**.")