src/streamlit_app.py

import os
os.environ['HF_HOME'] = '/tmp'

import os
import time
import streamlit as st
import pandas as pd
import io
import plotly.express as px
import zipfile
import json
from cryptography.fernet import Fernet
from streamlit_extras.stylable_container import stylable_container
from typing import Optional
from gliner import GLiNER
from comet_ml import Experiment
import hashlib

# Set up environment variables
os.environ['HF_HOME'] = '/tmp'

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 */
    }
    /* Sidebar background color */
    .css-1d36184 {
        background-color: #B2F2B2; /* A pale green for the sidebar */
        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 */
    }
    /* 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 */
        color: #000000;
    }
    /* Success box background and text color */
    .stAlert.st-success {
        background-color: #C8F0C8; /* A light green for the success box */
        color: #000000;
    }
    </style>
    """,
    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-six (36) 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"
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""")

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>
    '''
    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")

# --- 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", "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": ["Job_title", "Date", "Organization", "Role"],
    "Performance": ["Performance_score"],
    "Attendance": ["Leave_of_absence"],
    "Benefits": ["Retirement_plan", "Bonus", "Stock_options", "Health_insurance"],
    "Compensation": ["Pay_rate", "Annual_salary"],
    "Deductions": ["Tax", "Deductions"],
    "Recruitment & Sourcing": ["Interview_type", "Applicant", "Referral", "Job_board", "Recruiter"],
    "Legal & Compliance": ["Offer_letter", "Agreement"],
    "Professional_Development": ["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, 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}

# --- 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
    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)

# --- Results Section ---
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 = 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,
                        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 = 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.")

                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.")
                # Clear session state if no results found
                if 'df' in st.session_state:
                    del st.session_state.df

# --- 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 ---
@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()

qa_model = load_gliner_model()
st.subheader("Question-Answering", divider="violet")

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:
            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:
    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.")
    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)

        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:
                    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.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))

                    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",
                        )
                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}")