import os os.environ['HF_HOME'] = '/tmp' import time import streamlit as st import pandas as pd import io import plotly.express as px import zipfile import string 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 # --- Page Configuration and UI Elements --- st.set_page_config(layout="wide", page_title="Named Entity Recognition App") st.subheader("DataHarvest", divider="violet") st.link_button("by nlpblogs", "https://nlpblogs.com", type="tertiary") st.markdown(':rainbow[**Supported Languages: English**]') expander = st.expander("**Important notes**") expander.write("""**Named Entities:** This DataHarvest web app predicts nine (9) labels: "person", "country", "city", "organization", "date", "time", "cardinal", "money", "position" 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. **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 DataHarvest web app on your website. Feel free to adjust the width and height values to fit your page.") code = ''' ''' st.code(code, language="html") st.text("") st.text("") st.divider() st.subheader("🚀 Ready to build your own AI Web App?", divider="violet") st.link_button("AI Web App Builder", "https://nlpblogs.com/build-your-named-entity-recognition-app/", 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.") print("Warning: Comet ML environment variables are not set. Logging will be disabled.") # --- Label Definitions --- labels = ["person", "country", "city", "organization", "date", "time", "cardinal", "money", "position"] category_mapping = { "People": ["person", "organization", "position"], "Locations": ["country", "city"], "Time": ["date", "time"], "Numbers": ["money", "cardinal"] } # --- 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() reverse_category_mapping = {label: category for category, label_list in category_mapping.items() for label in label_list} # --- Session State Initialization --- if 'show_results' not in st.session_state: st.session_state.show_results = False if 'last_text' not in st.session_state: st.session_state.last_text = "" if 'results_df' not in st.session_state: st.session_state.results_df = pd.DataFrame() if 'elapsed_time' not in st.session_state: st.session_state.elapsed_time = 0.0 # --- Text Input and Clear Button --- word_limit = 200 text = st.text_area(f"Type or paste your text below (max {word_limit} words), and then press Ctrl + Enter", height=250, key='my_text_area') word_count = len(text.split()) st.markdown(f"**Word count:** {word_count}/{word_limit}") def clear_text(): """Clears the text area and hides results.""" st.session_state['my_text_area'] = "" st.session_state.show_results = False st.session_state.last_text = "" st.session_state.results_df = pd.DataFrame() st.session_state.elapsed_time = 0.0 def remove_punctuation(text): """Removes punctuation from a string.""" translator = str.maketrans('', '', string.punctuation) return text.translate(translator) st.button("Clear text", on_click=clear_text) # --- Results Section --- if st.button("Results"): if not text.strip(): st.warning("Please enter some text to extract entities.") st.session_state.show_results = False elif word_count > word_limit: st.warning(f"Your text exceeds the {word_limit} word limit. Please shorten it to continue.") st.session_state.show_results = False else: # Check if the text is different from the last time if text != st.session_state.last_text: st.session_state.show_results = True st.session_state.last_text = text start_time = time.time() with st.spinner("Extracting entities...", show_time=True): cleaned_text = remove_punctuation(text) entities = model.predict_entities(cleaned_text, labels) df = pd.DataFrame(entities) st.session_state.results_df = df 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) experiment.end() end_time = time.time() st.session_state.elapsed_time = end_time - start_time # If the text is the same, do nothing but keep results displayed else: st.session_state.show_results = True # Display results if the state variable is True if st.session_state.show_results: df = st.session_state.results_df if not df.empty: df['category'] = df['label'].map(reverse_category_mapping) st.subheader("Grouped Entities by Category", divider="violet") 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'] - **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="violet") 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="violet") 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="violet") 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="violet") 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'}) 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'], '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']}) 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: #8C8C8C; 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" ) st.text("") st.text("") st.info(f"Results processed in **{st.session_state.elapsed_time:.2f} seconds**.") else: st.warning("No entities were found in the provided text.")