app.py

import time
import streamlit as st
import pandas as pd
import io
from transformers import pipeline
import plotly.express as px
import zipfile
import re
import numpy as np
import json
import os
from comet_ml import Experiment

# --- Page Configuration ---
st.set_page_config(layout="wide", page_title="Named Entity Recognition App")

# Define the categories and their associated entity labels
ENTITY_LABELS_CATEGORIZED = {
    "Persons": ["PER"],
    "Locations": ["LOC"],
    "Organizations": ["ORG"],
    "Nationalities, Religious, Political Groups": ["NORP"],
    "Miscellaneous": ["MISC"],
}

# Create a mapping from each specific entity label to its category
LABEL_TO_CATEGORY_MAP = {
    label: category for category, labels in ENTITY_LABELS_CATEGORIZED.items() for label in labels
}

# --- Comet ML Configuration ---
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 = False
if COMET_API_KEY and COMET_WORKSPACE and COMET_PROJECT_NAME:
    comet_initialized = True
else:
    st.warning("Comet ML environment variables (COMET_API_KEY, COMET_WORKSPACE, COMET_PROJECT_NAME) not set. "
               "Comet ML logging will be skipped.")

@st.cache_resource
def load_ner_model():
    """
    Loads the pre-trained NER model ("UGARIT/grc-ner-bert") and caches it.
    """
    try:
        return pipeline(
            "token-classification",
            model="UGARIT/grc-ner-bert",
            aggregation_strategy="max",
            ignore_labels=["O"],
            stride=128
        )
    except Exception as e:
        st.error(f"Failed to load NER model. Please check your internet connection or model availability: {e}")
        st.stop()

# --- UI Elements ---
st.subheader("Free Ancient Greek Entity Finder", divider="orange")
st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary")
expander = st.expander("**Important notes on the Free Ancient Greek Entity Finder**")
expander.write('''
   **Named Entities:** This Free Ancient Greek Entity Finder predicts five
   (5) labels ("PER: person", "LOC: location", "ORG: organization", "NORP: Nationalities, Religious, Political Groups", "MISC:
   miscellaneous"). Results are presented in an easy-to-read table, 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 Ancient Greek text into the input box. Then, click the 'Results' button
   to extract and tag entities.
   
   **Language settings:** Please check and adjust the language settings in
   your computer, so the Ancient Greek characters are handled properly in your downloaded file.
   
   **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:
    container = st.container(border=True)
    container.write("**Named Entity Recognition (NER)** is the task of "
                    "extracting and tagging entities in text data. Entities can be persons, "
                    "organizations, locations, countries, products, events etc.")
    st.subheader("Related NER Web Apps", divider="orange")
    st.link_button("Multilingual PDF & DOCX Entity Finder",
                   "https://nlpblogs.com/shop/named-entity-recognition-ner/multilingual-pdf-docx-entity-finder/",
                   type="primary")



text_input = st.text_area("Type or paste your Ancient Greek text below, and then press Ctrl + Enter", key='my_text_area')
st.write("**Input text**: ", text_input)

def clear_text():
    st.session_state['my_text_area'] = ""

st.button("Clear text", on_click=clear_text)

st.divider()






# --- Results Button and Processing Logic ---
if st.button("Results"):
    start_time_overall = time.time()  # Start time for overall processing

    if not text_input.strip():
        st.warning("Please enter some text for analysis.")
        st.stop()

    experiment = None
    if comet_initialized:
        experiment = Experiment(
            api_key=COMET_API_KEY,
            workspace=COMET_WORKSPACE,
            project_name=COMET_PROJECT_NAME,
        )
        experiment.log_parameter("input_text", text_input)

    with st.spinner("Analyzing text...", show_time=True):
        model = load_ner_model()
        # Measure NER model processing time
        start_time_ner = time.time()
        text_entities = model(text_input)
        end_time_ner = time.time()
        ner_processing_time = end_time_ner - start_time_ner

        if experiment:
            experiment.log_metric("ner_processing_time_seconds", ner_processing_time)

        df = pd.DataFrame(text_entities)

        # --- Add 'category' column to the DataFrame based on the grouped labels ---
        df['category'] = df['entity_group'].map(LABEL_TO_CATEGORY_MAP)
        # Handle cases where an entity_group might not have a category
        df['category'] = df['category'].fillna('Uncategorized')

        if experiment:
            experiment.log_table("predicted_entities", df)

        # --- Display Results ---
        st.subheader("Extracted Entities", divider="orange")
        properties = {"border": "2px solid gray", "color": "blue", "font-size": "16px"}
        df_styled = df.style.set_properties(**properties)
        st.dataframe(df_styled, use_container_width=True)

        with st.expander("See Glossary of tags"):
            st.write('''
            '**word**': ['entity extracted from your text data']
            
            '**score**': ['accuracy score; how accurately a tag has been assigned to
            a given entity']
            
            '**entity_group**': ['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']
            
            '**category**': ['the broader category the entity belongs to']
            ''')

        st.subheader("Grouped entities", divider="orange")
        # Get unique categories and sort them for consistent tab order
        unique_categories = sorted(df['category'].unique())
        tabs_per_row = 5  # Adjust as needed for better layout
        # Loop through categories in chunks to create rows of tabs
        for i in range(0, len(unique_categories), tabs_per_row):
            current_row_categories = unique_categories[i : i + tabs_per_row]
            tabs = st.tabs(current_row_categories)
            for j, category in enumerate(current_row_categories):
                with tabs[j]:
                    df_filtered = df[df["category"] == category]
                    if not df_filtered.empty:
                        st.dataframe(df_filtered, use_container_width=True)
                    else:
                        st.info(f"No '{category}' entities found in the text.")
                        # Display an empty DataFrame for consistency if no entities are found
                        st.dataframe(pd.DataFrame({
                            'entity_group': [np.nan],
                            'score': [np.nan],
                            'word': [np.nan],
                            'start': [np.nan],
                            'end': [np.nan],
                            'category': [category]
                        }), hide_index=True)
        st.divider()

        # --- Visualizations ---
        st.subheader("Tree map", divider="orange")
        fig_treemap = px.treemap(df,
                                 path=[px.Constant("all"), 'category', 'entity_group', 'word'],
                                 values='score', color='category',
                                 color_discrete_map={
                                     'Persons': 'blue',
                                     'Locations': 'green',
                                     'Organizations': 'red',
                                     'Miscellaneous': 'purple',
                                     'Uncategorized': 'gray'
                                 })
        fig_treemap.update_layout(margin=dict(t=50, l=25, r=25, b=25))
        st.plotly_chart(fig_treemap)
        if experiment:
            experiment.log_figure(figure=fig_treemap, figure_name="entity_treemap")

        # Group by category and entity_group to get counts for pie and bar charts
        grouped_counts = df.groupby('category').size().reset_index(name='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.update_traces(textposition='inside', textinfo='percent+label')
            st.plotly_chart(fig_pie)
            if experiment:
                experiment.log_figure(figure=fig_pie, figure_name="category_pie_chart")

        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')
            st.plotly_chart(fig_bar)
            if experiment:
                experiment.log_figure(figure=fig_bar, figure_name="category_bar_chart")

        # --- Downloadable Content ---
        dfa = pd.DataFrame(
            data={
                'Column Name': ['word', 'entity_group', '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 zip file",
            data=buf.getvalue(),
            file_name="nlpblogs_ner_results.zip",
            mime="application/zip",
        )
        if experiment:
            experiment.log_asset_data(buf.getvalue(), file_name="nlpblogs_ner_results.zip",
                                      metadata={"type": "results_archive"})

    end_time_overall = time.time()
    elapsed_time_overall = end_time_overall - start_time_overall
    st.info(f"Results processed in **{elapsed_time_overall:.2f} seconds**.")

    if experiment:
        experiment.log_metric("overall_processing_time_seconds", elapsed_time_overall)
        experiment.end()