from flask import Flask, render_template_string, request, jsonify from flask_cors import CORS from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline import os import sys import threading import time app = Flask(__name__) CORS(app) # Model loading state (thread-safe) model_name = "OpenMed/privacy-filter-nemotron" classifier = None tokenizer = None model_loading = False model_error = None model_thread = None # Background model loading def load_model_async(): global classifier, tokenizer, model_loading, model_error model_loading = True print("="*60, flush=True) print("BACKGROUND: Loading OpenMed Privacy Filter model...", flush=True) print("="*60, flush=True) try: print(f"Loading tokenizer and model: {model_name}", flush=True) print("This may take 5-10 minutes on first run...", flush=True) # Use AutoModelForTokenClassification directly for better performance tokenizer = AutoTokenizer.from_pretrained( model_name, cache_dir="/app/.cache/huggingface" ) model = AutoModelForTokenClassification.from_pretrained( model_name, cache_dir="/app/.cache/huggingface" ) global classifier classifier = pipeline( task="token-classification", model=model, tokenizer=tokenizer, aggregation_strategy="first", #none simple first average max device=-1 ) print("✓ Model loaded successfully!", flush=True) model_error = None except Exception as e: model_error = str(e) print(f"✗ ERROR loading model: {e}", flush=True) import traceback traceback.print_exc() finally: model_loading = False # Start model loading in background model_thread = threading.Thread(target=load_model_async, daemon=True) model_thread.start() def escape_html(text): """Escape HTML special characters to prevent XSS""" return (text .replace("&", "&") .replace("<", "<") .replace(">", ">") .replace('"', """) .replace("'", "'")) def create_line_chunks(text, max_tokens=2048): """Split text into chunks that respect line boundaries. Groups lines together based on max token limit, never cutting mid-line. """ global tokenizer if tokenizer is None: return [(0, text, len(text.split()))] lines = text.split('\n') chunks = [] current_lines = [] current_token_count = 0 current_char_start = 0 for line in lines: line_tokens = tokenizer(line, add_special_tokens=False)['input_ids'] line_token_count = len(line_tokens) # If this single line exceeds max_tokens, we have to include it anyway if current_token_count + line_token_count > max_tokens and current_lines: # Save current chunk chunk_text = '\n'.join(current_lines) chunks.append((current_char_start, chunk_text, current_token_count)) # Start new chunk with this line current_lines = [line] current_token_count = line_token_count current_char_start = text.find(line, current_char_start + len('\n'.join(current_lines[:-1])) if current_lines[:-1] else 0) else: current_lines.append(line) current_token_count += line_token_count # Add final chunk if current_lines: chunk_text = '\n'.join(current_lines) chunks.append((current_char_start, chunk_text, current_token_count)) return chunks def merge_adjacent_entities(entities): """Merge adjacent entities of the same type that are likely from tokenization splits.""" if not entities: return entities # Sort by start position sorted_entities = sorted(entities, key=lambda x: x.get('start', 0)) merged = [] i = 0 while i < len(sorted_entities): current = sorted_entities[i] current_label = current.get('entity_group') or current.get('entity', 'unknown') current_end = current.get('end', 0) current_text = current.get('word', '') current_score = current.get('score', 0) # Look ahead for adjacent same-type entities j = i + 1 while j < len(sorted_entities): next_entity = sorted_entities[j] next_label = next_entity.get('entity_group') or next_entity.get('entity', 'unknown') next_start = next_entity.get('start', 0) # Check if same label and adjacent (or overlapping/nearby) if next_label == current_label and next_start <= current_end + 5: # Merge next_end = next_entity.get('end', 0) next_text = next_entity.get('word', '') next_score = next_entity.get('score', 0) # Combine text (remove overlap if any) if next_start <= current_end: current_text = current_text[:next_start - current.get('start', 0)] + next_text else: current_text = current_text + ' ' + next_text current_end = max(current_end, next_end) current_score = max(current_score, next_score) # Use highest score j += 1 else: break merged.append({ 'entity_group': current_label, 'entity': current_label, 'word': current_text, 'start': current.get('start', 0), 'end': current_end, 'score': current_score }) i = j return merged # HTML Template with proper loading states HTML_TEMPLATE = ''' OpenMed Privacy Filter + Nemotron - PII Detection Demo

OpenMed Privacy Filter + Nemotron

PII Detection & Masking Demo using Flask

Waiting for server to start...

Detects 55 Types of PII:

Max chunk: 128000 (model limit)

Results

''' @app.route('/') def index(): return render_template_string(HTML_TEMPLATE) @app.route('/health') def health(): """Health check with model loading status""" global classifier, model_loading, model_error, model_thread if classifier is not None: return jsonify({ 'status': 'healthy', 'model_loaded': True, 'model_loading': False }) elif model_loading: return jsonify({ 'status': 'loading', 'model_loaded': False, 'model_loading': True, 'message': 'Model is still loading, please wait...' }) else: return jsonify({ 'status': 'unhealthy', 'model_loaded': False, 'model_loading': False, 'error': model_error or 'Model loading failed or thread terminated unexpectedly' }), 503 @app.route('/tokenize', methods=['POST']) def tokenize(): """Count tokens in the provided text.""" global tokenizer if tokenizer is None: return jsonify({'success': False, 'error': 'Tokenizer not loaded yet'}), 503 try: data = request.get_json() if not data: return jsonify({'success': False, 'error': 'No JSON data'}), 400 text = data.get('text', '') if not text: return jsonify({'token_count': 0}) # Tokenize and count tokens = tokenizer(text, add_special_tokens=False) token_count = int(len(tokens['input_ids'])) return jsonify({ 'success': True, 'token_count': token_count }) except Exception as e: return jsonify({'success': False, 'error': str(e)}), 500 @app.route('/analyze', methods=['POST', 'OPTIONS']) def analyze(): if request.method == 'OPTIONS': return '', 204 global classifier, tokenizer if classifier is None: return jsonify({ 'success': False, 'error': f'Model not yet loaded. Please wait and refresh in a few minutes.' }), 503 try: data = request.get_json() if not data: return jsonify({'success': False, 'error': 'No JSON data received'}), 400 text = data.get('text', '') chunk_size = data.get('chunk_size', 4096) if not text.strip(): return jsonify({'success': True, 'entities': [], 'entity_count': 0, 'chunks_processed': 0}) # Create line-based chunks chunks = create_line_chunks(text, max_tokens=chunk_size) all_entities = [] # Process each chunk for char_offset, chunk_text, chunk_tokens in chunks: chunk_results = classifier(chunk_text) # Adjust entity positions by chunk offset for entity in chunk_results: entity['start'] = entity.get('start', 0) + char_offset entity['end'] = entity.get('end', 0) + char_offset all_entities.append(entity) # Merge adjacent entities merged_entities = merge_adjacent_entities(all_entities) # Format output entities = [] for entity in merged_entities: entities.append({ 'label': entity.get('entity_group', entity.get('entity', 'unknown')), 'text': entity.get('word', ''), 'start': entity.get('start', 0), 'end': entity.get('end', 0), 'score': float(entity.get('score', 0)) }) return jsonify({ 'success': True, 'entities': entities, 'entity_count': len(entities), 'chunks_processed': len(chunks) }) except Exception as e: print(f"Error during analysis: {e}", flush=True) import traceback traceback.print_exc() return jsonify({ 'success': False, 'error': str(e) }), 500 if __name__ == '__main__': port = int(os.environ.get('PORT', 7860)) app.run(host='0.0.0.0', port=port, debug=False, threaded=True)