import os import csv import streamlit as st import polars as pl from io import BytesIO, StringIO from gliner import GLiNER from gliner_file import run_ner import time import torch import platform from typing import List from streamlit_tags import st_tags # Importing the st_tags component for labels # Streamlit page configuration st.set_page_config( page_title="GLiNER", page_icon="🔥", layout="wide", initial_sidebar_state="expanded" ) # Function to load data from the uploaded file @st.cache_data def load_data(file): """ Loads an uploaded CSV or Excel file with resilient detection of delimiters and types. """ with st.spinner("Loading data, please wait..."): try: _, file_ext = os.path.splitext(file.name) if file_ext.lower() in [".xls", ".xlsx"]: return load_excel(file) elif file_ext.lower() == ".csv": return load_csv(file) else: raise ValueError("Unsupported file format. Please upload a CSV or Excel file.") except Exception as e: st.error("Error loading data:") st.error(str(e)) return None def load_excel(file): """ Loads an Excel file using `BytesIO` and `polars` for reduced latency. """ try: file_bytes = BytesIO(file.read()) df = pl.read_excel(file_bytes, read_options={"ignore_errors": True}) return df except Exception as e: raise ValueError(f"Error reading the Excel file: {str(e)}") def load_csv(file): """ Loads a CSV file by detecting the delimiter and using the quote character to handle internal delimiters. """ try: file.seek(0) # Reset file pointer to ensure reading from the beginning raw_data = file.read() try: file_content = raw_data.decode('utf-8') except UnicodeDecodeError: try: file_content = raw_data.decode('latin1') except UnicodeDecodeError: raise ValueError("Unable to decode the file. Ensure it is encoded in UTF-8 or Latin-1.") delimiters = [",", ";", "|", "\t", " "] for delimiter in delimiters: try: df = pl.read_csv( StringIO(file_content), separator=delimiter, quote_char='"', try_parse_dates=True, ignore_errors=True, truncate_ragged_lines=True ) return df except Exception: continue raise ValueError("Unable to load the file with common delimiters.") except Exception as e: raise ValueError(f"Error reading the CSV file: {str(e)}") @st.cache_resource def load_model(): """ Loads the GLiNER model into memory to avoid multiple reloads. """ try: gpu_available = torch.cuda.is_available() with st.spinner("Loading the GLiNER model... Please wait."): device = torch.device("cuda" if gpu_available else "cpu") model = GLiNER.from_pretrained( "urchade/gliner_multi-v2.1" ).to(device) model.eval() if gpu_available: device_name = torch.cuda.get_device_name(0) st.success(f"GPU detected: {device_name}. Model loaded on GPU.") else: cpu_name = platform.processor() st.warning(f"No GPU detected. Using CPU: {cpu_name}") return model except Exception as e: st.error("Error loading the model:") st.error(str(e)) return None def perform_ner(filtered_df, selected_column, labels_list, threshold): """ Executes named entity recognition (NER) on the filtered data. """ try: texts_to_analyze = filtered_df[selected_column].to_list() total_rows = len(texts_to_analyze) ner_results_list = [] progress_bar = st.progress(0) progress_text = st.empty() start_time = time.time() for index, text in enumerate(texts_to_analyze, 1): if st.session_state.stop_processing: progress_text.text("Processing stopped by user.") break ner_results = run_ner( st.session_state.gliner_model, [text], labels_list, threshold=threshold ) ner_results_list.append(ner_results) progress = index / total_rows elapsed_time = time.time() - start_time progress_bar.progress(progress) progress_text.text(f"Progress: {index}/{total_rows} - {progress * 100:.0f}% (Elapsed time: {elapsed_time:.2f}s)") for label in labels_list: extracted_entities = [] for entities in ner_results_list: texts = [entity["text"] for entity in entities[0] if entity["label"] == label] concatenated_texts = ", ".join(texts) if texts else "" extracted_entities.append(concatenated_texts) filtered_df = filtered_df.with_columns(pl.Series(name=label, values=extracted_entities)) end_time = time.time() st.success(f"Processing completed in {end_time - start_time:.2f} seconds.") return filtered_df except Exception as e: st.error(f"Error during NER processing: {str(e)}") return filtered_df def main(): st.title("Use NER with GliNER on your data file") st.markdown("Prototype v0.1") st.write(""" This application performs named entity recognition (NER) on your text data using GLiNER. **Instructions:** 1. Upload a CSV or Excel file. 2. Select the column containing the text to analyze. 3. Filter the data if necessary. 4. Enter the NER labels you wish to detect. 5. Click "Start NER" to begin processing. """) if "stop_processing" not in st.session_state: st.session_state.stop_processing = False if "threshold" not in st.session_state: st.session_state.threshold = 0.4 if "labels_list" not in st.session_state: st.session_state.labels_list = [] st.session_state.gliner_model = load_model() if st.session_state.gliner_model is None: return uploaded_file = st.sidebar.file_uploader("Choose a file (CSV or Excel)") if uploaded_file is None: st.warning("Please upload a file to continue.") return df = load_data(uploaded_file) if df is None: return selected_column = st.selectbox("Select the column containing the text:", df.columns) filter_text = st.text_input("Filter the column by text", "") if filter_text: filtered_df = df.filter(pl.col(selected_column).str.contains(f"(?i).*{filter_text}.*")) else: filtered_df = df st.write("Filtered data preview:") rows_per_page = 100 total_rows = len(filtered_df) total_pages = (total_rows - 1) // rows_per_page + 1 if "current_page" not in st.session_state: st.session_state.current_page = 1 def update_page(new_page): st.session_state.current_page = new_page col1, col2, col3, col4, col5 = st.columns(5) with col1: first = st.button("⏮️ First") with col2: previous = st.button("⬅️ Previous") with col3: pass with col4: next = st.button("Next ➡️") with col5: last = st.button("Last ⏭️") if first: update_page(1) elif previous: if st.session_state.current_page > 1: update_page(st.session_state.current_page - 1) elif next: if st.session_state.current_page < total_pages: update_page(st.session_state.current_page + 1) elif last: update_page(total_pages) with col3: st.markdown(f"Page **{st.session_state.current_page}** of **{total_pages}**") start_idx = (st.session_state.current_page - 1) * rows_per_page end_idx = min(start_idx + rows_per_page, total_rows) if not filtered_df.is_empty(): current_page_data = filtered_df.slice(start_idx, end_idx - start_idx) st.write(f"Displaying {start_idx + 1} to {end_idx} of {total_rows} rows") st.dataframe(current_page_data.to_pandas(), use_container_width=True) else: st.warning("The filtered DataFrame is empty. Please check your filters.") st.slider("Set confidence threshold", 0.0, 1.0, st.session_state.threshold, 0.01, key="threshold") st.session_state.labels_list = st_tags( label="Enter the NER labels to detect", text="Add more labels as needed", value=st.session_state.labels_list, key="1" ) col1, col2 = st.columns(2) with col1: start_button = st.button("Start NER") with col2: stop_button = st.button("Stop") if start_button: st.session_state.stop_processing = False if not st.session_state.labels_list: st.warning("Please enter labels for NER.") else: updated_df = perform_ner(filtered_df, selected_column, st.session_state.labels_list, st.session_state.threshold) st.write("**NER Results:**") st.dataframe(updated_df.to_pandas(), use_container_width=True) def to_excel(df): output = BytesIO() df.write_excel(output) return output.getvalue() def to_csv(df): return df.write_csv().encode('utf-8') download_col1, download_col2 = st.columns(2) with download_col1: st.download_button( label="📥 Download as Excel", data=to_excel(updated_df), file_name="ner_results.xlsx"#, #mime="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", ) with download_col2: st.download_button( label="📥 Download as CSV", data=to_csv(updated_df), file_name="ner_results.csv"#, #mime="text/csv", ) if stop_button: st.session_state.stop_processing = True st.warning("Processing stopped by user.") if __name__ == "__main__": main()