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 # 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: # Load the file into BytesIO for faster reading file_bytes = BytesIO(file.read()) # Load the Excel file using `polars` 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 decoding as UTF-8, else as Latin-1 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.") # List of common delimiters delimiters = [",", ";", "|", "\t", " "] # Try each delimiter until one works for delimiter in delimiters: try: # Read CSV with current delimiter and handle quoted fields df = pl.read_csv( StringIO(file_content), separator=delimiter, quote_char='"', # Handle internal delimiters with quotes try_parse_dates=True, ignore_errors=True, # Ignore errors for invalid values truncate_ragged_lines=True ) # Return the DataFrame if loading succeeds return df except Exception: continue # Move to the next delimiter in case of error # If no delimiter worked raise ValueError("Unable to load the file with common delimiters.") except Exception as e: raise ValueError(f"Error reading the CSV file: {str(e)}") # Function to load the GLiNER model @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 # Function to perform NER and update the user interface 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 = [] # Initialize progress bar and text progress_bar = st.progress(0) progress_text = st.empty() start_time = time.time() # Process each row individually to keep progress updates responsive 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) # Update progress bar and text after each row 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)") # Add NER results to the DataFrame 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 # Main function to run the Streamlit application def main(): st.title("Use NER with GliNER on your data file") st.markdown("Prototype v0.1") # User instructions 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. """) # Initializing session state variables 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 = [] # Load the model st.session_state.gliner_model = load_model() if st.session_state.gliner_model is None: return # File upload 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 # Loading data df = load_data(uploaded_file) if df is None: return # Column selection selected_column = st.selectbox("Select the column containing the text:", df.columns) # Data filtering 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 rows_per_page = 100 # Calculate total rows and pages total_rows = len(filtered_df) total_pages = (total_rows - 1) // rows_per_page + 1 # Initialize current page in session_state if "current_page" not in st.session_state: st.session_state.current_page = 1 # Function to update page def update_page(new_page): st.session_state.current_page = new_page # Pagination buttons col1, col2, col3, col4, col5 = st.columns(5) with col1: first = st.button("⏮️ First") with col2: previous = st.button("⬅️ Previous") with col3: pass # Page number display will be done after with col4: next = st.button("Next ➡️") with col5: last = st.button("Last ⏭️") # Button clicks management 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) # Now display the page number after updating with col3: st.markdown(f"Page **{st.session_state.current_page}** of **{total_pages}**") # Calculate indices for pagination start_idx = (st.session_state.current_page - 1) * rows_per_page end_idx = min(start_idx + rows_per_page, total_rows) # Check if the filtered DataFrame is empty if not filtered_df.is_empty(): # Retrieve current page data 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.") # Confidence threshold slider st.slider("Set confidence threshold", 0.0, 1.0, st.session_state.threshold, 0.01, key="threshold") # Buttons to start and stop NER 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: # Run NER 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) # Function to convert DataFrame to Excel def to_excel(df): output = BytesIO() df.write_excel(output) return output.getvalue() # Function to convert DataFrame to CSV def to_csv(df): return df.write_csv().encode('utf-8') # Download buttons for results 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()