profanity_detector.py

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForSeq2SeqLM
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
import whisper
import gradio as gr
import re
import pandas as pd
import numpy as np
import os
import time
import logging
import threading
import queue
from scipy.io.wavfile import write as write_wav
from html import escape
import traceback
import spaces # Required for Hugging Face ZeroGPU compatibility

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[logging.StreamHandler()]
)
logger = logging.getLogger('profanity_detector')

# Detect if we're running in a ZeroGPU environment
IS_ZEROGPU = os.environ.get("SPACE_RUNTIME_STATELESS", "0") == "1"
if os.environ.get("SPACES_ZERO_GPU") is not None:
    IS_ZEROGPU = True

# Define device strategy that works in both environments
if IS_ZEROGPU:
    # In ZeroGPU: always initialize on CPU, will use GPU only in @spaces.GPU functions
    device = torch.device("cpu")
    logger.info("ZeroGPU environment detected. Using CPU for initial loading.")
else:
    # For local runs: use CUDA if available
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Local environment. Using device: {device}")

# Global variables for models
profanity_model = None
profanity_tokenizer = None
t5_model = None
t5_tokenizer = None
whisper_model = None
tts_processor = None
tts_model = None
vocoder = None
models_loaded = False

# Default speaker embeddings for TTS
speaker_embeddings = None

# Queue for real-time audio processing
audio_queue = queue.Queue()
processing_active = False

# Model loading with int8 quantization
def load_models():
    global profanity_model, profanity_tokenizer, t5_model, t5_tokenizer, whisper_model
    global tts_processor, tts_model, vocoder, speaker_embeddings, models_loaded
    
    try:
        logger.info("Loading profanity detection model...")
        PROFANITY_MODEL = "parsawar/profanity_model_3.1"
        profanity_tokenizer = AutoTokenizer.from_pretrained(PROFANITY_MODEL)
        
        # Load model without moving to CUDA directly
        profanity_model = AutoModelForSequenceClassification.from_pretrained(
            PROFANITY_MODEL,
            device_map=None,  # Stay on CPU for now
            low_cpu_mem_usage=True
        )
        
        # Only move to device if NOT in ZeroGPU mode
        if not IS_ZEROGPU and torch.cuda.is_available():
            profanity_model = profanity_model.to(device)
            try:
                profanity_model = profanity_model.half()
                logger.info("Successfully converted profanity model to half precision")
            except Exception as e:
                logger.warning(f"Could not convert to half precision: {str(e)}")
        
        logger.info("Loading detoxification model...")
        T5_MODEL = "s-nlp/t5-paranmt-detox"
        t5_tokenizer = AutoTokenizer.from_pretrained(T5_MODEL)
        
        t5_model = AutoModelForSeq2SeqLM.from_pretrained(
            T5_MODEL,
            device_map=None,  # Stay on CPU for now
            low_cpu_mem_usage=True
        )
        
        # Only move to device if NOT in ZeroGPU mode
        if not IS_ZEROGPU and torch.cuda.is_available():
            t5_model = t5_model.to(device)
            try:
                t5_model = t5_model.half()
                logger.info("Successfully converted T5 model to half precision")
            except Exception as e:
                logger.warning(f"Could not convert to half precision: {str(e)}")
        
        logger.info("Loading Whisper speech-to-text model...")
        # Always load on CPU in ZeroGPU mode
        #whisper_model = whisper.load_model("medium" if IS_ZEROGPU else "large", device="cpu")
        whisper_model = whisper.load_model("large-v2", device="cpu")
        
        # Only move to device if NOT in ZeroGPU mode
        if not IS_ZEROGPU and torch.cuda.is_available():
            whisper_model = whisper_model.to(device)
            
        logger.info("Loading Text-to-Speech model...")
        TTS_MODEL = "microsoft/speecht5_tts"
        tts_processor = SpeechT5Processor.from_pretrained(TTS_MODEL)
        
        tts_model = SpeechT5ForTextToSpeech.from_pretrained(
            TTS_MODEL,
            device_map=None,  # Stay on CPU for now
            low_cpu_mem_usage=True
        )
        
        vocoder = SpeechT5HifiGan.from_pretrained(
            "microsoft/speecht5_hifigan",
            device_map=None,  # Stay on CPU for now
            low_cpu_mem_usage=True
        )
        
        # Only move to device if NOT in ZeroGPU mode
        if not IS_ZEROGPU and torch.cuda.is_available():
            tts_model = tts_model.to(device)
            vocoder = vocoder.to(device)
        
        # Speaker embeddings - always on CPU for ZeroGPU
        speaker_embeddings = torch.zeros((1, 512))
        # Only move to device if NOT in ZeroGPU mode
        if not IS_ZEROGPU and torch.cuda.is_available():
            speaker_embeddings = speaker_embeddings.to(device)
            
        models_loaded = True
        logger.info("All models loaded successfully.")
        
        return "Models loaded successfully."
    except Exception as e:
        error_msg = f"Error loading models: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return error_msg

# ZeroGPU decorator: Requests GPU resources when function is called and releases them when completed.
# This enables efficient GPU sharing in Hugging Face Spaces while having no effect in local environments.
@spaces.GPU
def detect_profanity(text: str, threshold: float = 0.5):
    """
    Detect profanity in text with adjustable threshold
    
    Args:
        text: The input text to analyze
        threshold: Profanity detection threshold (0.0-1.0)
        
    Returns:
        Dictionary with analysis results
    """
    if not models_loaded:
        return {"error": "Models not loaded yet. Please wait."}
    
    try:
        # Detect profanity and score
        inputs = profanity_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        
        # In ZeroGPU, move to GPU here inside the spaces.GPU function
        # For local environments, it might already be on the correct device
        current_device = device
        if IS_ZEROGPU and torch.cuda.is_available():
            current_device = torch.device("cuda")
            inputs = inputs.to(current_device)
            # Only in ZeroGPU mode, we need to move the model to GPU inside the function
            profanity_model.to(current_device)
        elif torch.cuda.is_available():  # Local environment with CUDA
            inputs = inputs.to(current_device)
            
        with torch.no_grad():
            outputs = profanity_model(**inputs).logits
        score = torch.nn.functional.softmax(outputs, dim=1)[0][1].item()

        # Identify specific profane words
        words = re.findall(r'\b\w+\b', text)
        profane_words = []
        word_scores = {}
        
        if score > threshold:
            for word in words:
                if len(word) < 2:  # Skip very short words
                    continue
                    
                word_inputs = profanity_tokenizer(word, return_tensors="pt", truncation=True, max_length=512)
                if torch.cuda.is_available():
                    word_inputs = word_inputs.to(current_device)
                    
                with torch.no_grad():
                    word_outputs = profanity_model(**word_inputs).logits
                word_score = torch.nn.functional.softmax(word_outputs, dim=1)[0][1].item()
                word_scores[word] = word_score
                
                if word_score > threshold:
                    profane_words.append(word.lower())

        # Move model back to CPU if in ZeroGPU mode - to free GPU memory
        if IS_ZEROGPU and torch.cuda.is_available():
            profanity_model.to(torch.device("cpu"))

        # Create highlighted version of the text
        highlighted_text = create_highlighted_text(text, profane_words)

        return {
            "text": text, 
            "score": score, 
            "profanity": score > threshold, 
            "profane_words": profane_words,
            "highlighted_text": highlighted_text,
            "word_scores": word_scores
        }
    except Exception as e:
        error_msg = f"Error in profanity detection: {str(e)}"
        logger.error(error_msg)
        # Make sure model is on CPU if in ZeroGPU mode - to free GPU memory
        if IS_ZEROGPU and torch.cuda.is_available():
            try:
                profanity_model.to(torch.device("cpu"))
            except:
                pass
        return {"error": error_msg, "text": text, "score": 0, "profanity": False}

def create_highlighted_text(text, profane_words):
    """
    Create HTML-formatted text with profane words highlighted
    """
    if not profane_words:
        return escape(text)
        
    # Create a regex pattern matching any of the profane words (case insensitive)
    pattern = r'\b(' + '|'.join(re.escape(word) for word in profane_words) + r')\b'
    
    # Replace occurrences with highlighted versions
    def highlight_match(match):
        return f'<span style="background-color: rgba(255, 0, 0, 0.3); padding: 0px 2px; border-radius: 3px;">{match.group(0)}</span>'
    
    highlighted = re.sub(pattern, highlight_match, text, flags=re.IGNORECASE)
    return highlighted

@spaces.GPU
def rephrase_profanity(text):
    """
    Rephrase text containing profanity
    """
    if not models_loaded:
        return "Models not loaded yet. Please wait."
        
    try:
        # Rephrase using the detoxification model
        inputs = t5_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        
        # In ZeroGPU, move to GPU here inside the spaces.GPU function
        current_device = device
        if IS_ZEROGPU and torch.cuda.is_available():
            current_device = torch.device("cuda")
            inputs = inputs.to(current_device)
            # Only in ZeroGPU mode, we need to move the model to GPU inside the function
            t5_model.to(current_device)
        elif torch.cuda.is_available():  # Local environment with CUDA
            inputs = inputs.to(current_device)
        
        # Use more conservative generation settings with error handling
        try:
            outputs = t5_model.generate(
                **inputs,
                max_length=512,
                num_beams=4,         # Reduced from 5 to be more memory-efficient
                early_stopping=True,
                no_repeat_ngram_size=2,
                length_penalty=1.0
            )
            rephrased_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Verify the output is reasonable
            if not rephrased_text or len(rephrased_text) < 3:
                logger.warning(f"T5 model produced unusable output: '{rephrased_text}'")
                return text  # Return original if output is too short
                
            # Move model back to CPU if in ZeroGPU mode - to free GPU memory
            if IS_ZEROGPU and torch.cuda.is_available():
                t5_model.to(torch.device("cpu"))
                
            return rephrased_text.strip()
            
        except RuntimeError as e:
            # Handle potential CUDA out of memory error
            if "CUDA out of memory" in str(e):
                logger.warning("CUDA out of memory in T5 model. Trying with smaller beam size...")
                # Try again with smaller beam size
                outputs = t5_model.generate(
                    **inputs,
                    max_length=512,
                    num_beams=2,  # Use smaller beam size
                    early_stopping=True
                )
                rephrased_text = t5_tokenizer.decode(outputs[0], skip_special_tokens=True)
                
                # Move model back to CPU if in ZeroGPU mode - to free GPU memory
                if IS_ZEROGPU and torch.cuda.is_available():
                    t5_model.to(torch.device("cpu"))
                    
                return rephrased_text.strip()
            else:
                raise e  # Re-raise if it's not a memory issue
                
    except Exception as e:
        error_msg = f"Error in rephrasing: {str(e)}"
        logger.error(error_msg)
        # Make sure model is on CPU if in ZeroGPU mode - to free GPU memory
        if IS_ZEROGPU and torch.cuda.is_available():
            try:
                t5_model.to(torch.device("cpu"))
            except:
                pass
        return text  # Return original text if rephrasing fails

@spaces.GPU
def text_to_speech(text):
    """
    Convert text to speech using SpeechT5
    """
    if not models_loaded:
        return None
        
    try:
        # Create a temporary file path to save the audio
        temp_file = f"temp_tts_output_{int(time.time())}.wav"
        
        # Process the text input
        inputs = tts_processor(text=text, return_tensors="pt")
        
        # In ZeroGPU, move to GPU here inside the spaces.GPU function
        current_device = device
        if IS_ZEROGPU and torch.cuda.is_available():
            current_device = torch.device("cuda")
            inputs = inputs.to(current_device)
            # Only in ZeroGPU mode, we need to move the models to GPU inside the function
            tts_model.to(current_device)
            vocoder.to(current_device)
            speaker_embeddings_local = speaker_embeddings.to(current_device)
        elif torch.cuda.is_available():  # Local environment with CUDA
            inputs = inputs.to(current_device)
            speaker_embeddings_local = speaker_embeddings
        else:
            speaker_embeddings_local = speaker_embeddings
        
        # Generate speech with a fixed speaker embedding
        speech = tts_model.generate_speech(
            inputs["input_ids"], 
            speaker_embeddings_local, 
            vocoder=vocoder
        )
        
        # Convert from PyTorch tensor to NumPy array
        speech_np = speech.cpu().numpy()
        
        # Move models back to CPU if in ZeroGPU mode - to free GPU memory
        if IS_ZEROGPU and torch.cuda.is_available():
            tts_model.to(torch.device("cpu"))
            vocoder.to(torch.device("cpu"))
        
        # Save as WAV file (sampling rate is 16kHz for SpeechT5)
        write_wav(temp_file, 16000, speech_np)
        
        return temp_file
    except Exception as e:
        error_msg = f"Error in text-to-speech conversion: {str(e)}"
        logger.error(error_msg)
        # Make sure models are on CPU if in ZeroGPU mode - to free GPU memory
        if IS_ZEROGPU and torch.cuda.is_available():
            try:
                tts_model.to(torch.device("cpu"))
                vocoder.to(torch.device("cpu"))
            except:
                pass
        return None

def text_analysis(input_text, threshold=0.5):
    """
    Analyze text for profanity with adjustable threshold
    """
    if not models_loaded:
        return "Models not loaded yet. Please wait for initialization to complete.", None, None
        
    try:
        # Detect profanity with the given threshold
        result = detect_profanity(input_text, threshold=threshold)
        
        # Handle error case
        if "error" in result:
            return result["error"], None, None
            
        # Process results
        if result["profanity"]:
            clean_text = rephrase_profanity(input_text)
            profane_words_str = ", ".join(result["profane_words"])
            
            toxicity_score = result["score"]
            
            classification = (
                "Severe Toxicity" if toxicity_score >= 0.7 else
                "Moderate Toxicity" if toxicity_score >= 0.5 else
                "Mild Toxicity" if toxicity_score >= 0.35 else
                "Minimal Toxicity" if toxicity_score >= 0.2 else
                "No Toxicity"
            )
            
            # Generate audio for the rephrased text
            audio_output = text_to_speech(clean_text)
            
            return (
                f"Profanity Score: {result['score']:.4f}\n\n"
                f"Profane: {result['profanity']}\n"
                f"Classification: {classification}\n"
                f"Detected Profane Words: {profane_words_str}\n\n"
                f"Reworded: {clean_text}"
            ), result["highlighted_text"], audio_output
        else:
            # If no profanity detected, just convert the original text to speech
            audio_output = text_to_speech(input_text)
            
            return (
                f"Profanity Score: {result['score']:.4f}\n"
                f"Profane: {result['profanity']}\n"
                f"Classification: No Toxicity"
            ), None, audio_output
    except Exception as e:
        error_msg = f"Error in text analysis: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return error_msg, None, None

# ZeroGPU decorator with custom duration: Allocates GPU for up to 120 seconds to handle longer audio processing.
# Longer durations ensure processing isn't cut off, while shorter durations improve queue priority.
@spaces.GPU(duration=120)
def analyze_audio(audio_path, threshold=0.5):
    """
    Analyze audio for profanity with adjustable threshold
    """
    if not models_loaded:
        return "Models not loaded yet. Please wait for initialization to complete.", None, None
        
    if not audio_path:
        return "No audio provided.", None, None
        
    try:
        # In ZeroGPU mode, models need to be moved to GPU
        if IS_ZEROGPU and torch.cuda.is_available():
            current_device = torch.device("cuda")
            whisper_model.to(current_device)
        
        # Transcribe audio
        result = whisper_model.transcribe(audio_path, fp16=torch.cuda.is_available())
        text = result["text"]
        
        # Move whisper model back to CPU if in ZeroGPU mode
        if IS_ZEROGPU and torch.cuda.is_available():
            whisper_model.to(torch.device("cpu"))
        
        # Detect profanity with user-defined threshold
        analysis = detect_profanity(text, threshold=threshold)
        
        # Handle error case
        if "error" in analysis:
            return f"Error during analysis: {analysis['error']}\nTranscription: {text}", None, None

        if analysis["profanity"]:
            clean_text = rephrase_profanity(text)
        else:
            clean_text = text

        # Generate audio for the rephrased text
        audio_output = text_to_speech(clean_text)
        
        return (
            f"Transcription: {text}\n\n"
            f"Profanity Score: {analysis['score']:.4f}\n"
            f"Profane: {'Yes' if analysis['profanity'] else 'No'}\n"
            f"Classification: {'Severe Toxicity' if analysis['score'] >= 0.7 else 'Moderate Toxicity' if analysis['score'] >= 0.5 else 'Mild Toxicity' if analysis['score'] >= 0.35 else 'Minimal Toxicity' if analysis['score'] >= 0.2 else 'No Toxicity'}\n"
            f"Profane Words: {', '.join(analysis['profane_words']) if analysis['profanity'] else 'None'}\n\n"
            f"Reworded: {clean_text}"
        ), analysis["highlighted_text"] if analysis["profanity"] else None, audio_output
    except Exception as e:
        error_msg = f"Error in audio analysis: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        # Make sure models are on CPU if in ZeroGPU mode
        if IS_ZEROGPU and torch.cuda.is_available():
            try:
                whisper_model.to(torch.device("cpu"))
            except:
                pass
        return error_msg, None, None

# Global variables to store streaming results
stream_results = {
    "transcript": "",
    "profanity_info": "",
    "clean_text": "",
    "audio_output": None
}

@spaces.GPU
def process_stream_chunk(audio_chunk):
    """Process an audio chunk from the streaming interface"""
    global stream_results, processing_active
    
    if not processing_active or not models_loaded:
        return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
    
    try:
        # The format of audio_chunk from Gradio streaming can vary
        # It can be: (numpy_array, sample_rate), (filepath, sample_rate, numpy_array) or just numpy_array
        # Let's handle all possible cases
        
        if audio_chunk is None:
            # No audio received
            return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
        
        # Different Gradio versions return different formats
        temp_file = None
        
        if isinstance(audio_chunk, tuple):
            if len(audio_chunk) == 2:
                # Format: (numpy_array, sample_rate)
                samples, sample_rate = audio_chunk
                temp_file = f"temp_stream_{int(time.time())}.wav"
                write_wav(temp_file, sample_rate, samples)
            elif len(audio_chunk) == 3:
                # Format: (filepath, sample_rate, numpy_array)
                filepath, sample_rate, samples = audio_chunk
                # Use the provided filepath if it exists
                if os.path.exists(filepath):
                    temp_file = filepath
                else:
                    # Create our own file
                    temp_file = f"temp_stream_{int(time.time())}.wav"
                    write_wav(temp_file, sample_rate, samples)
        elif isinstance(audio_chunk, np.ndarray):
            # Just a numpy array, assume sample rate of 16000 for Whisper
            samples = audio_chunk
            sample_rate = 16000
            temp_file = f"temp_stream_{int(time.time())}.wav"
            write_wav(temp_file, sample_rate, samples)
        elif isinstance(audio_chunk, str) and os.path.exists(audio_chunk):
            # It's a filepath
            temp_file = audio_chunk
        else:
            # Unknown format
            stream_results["profanity_info"] = f"Error: Unknown audio format: {type(audio_chunk)}"
            return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
        
        # Make sure we have a valid file to process
        if not temp_file or not os.path.exists(temp_file):
            stream_results["profanity_info"] = "Error: Failed to create audio file for processing"
            return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
            
        # In ZeroGPU mode, move whisper model to GPU
        if IS_ZEROGPU and torch.cuda.is_available():
            current_device = torch.device("cuda")
            whisper_model.to(current_device)
            
        # Process with Whisper
        result = whisper_model.transcribe(temp_file, fp16=torch.cuda.is_available())
        transcript = result["text"].strip()
        
        # Move whisper model back to CPU if in ZeroGPU mode
        if IS_ZEROGPU and torch.cuda.is_available():
            whisper_model.to(torch.device("cpu"))
        
        # Skip processing if transcript is empty
        if not transcript:
            # Clean up temp file if we created it
            if temp_file and temp_file.startswith("temp_stream_") and os.path.exists(temp_file):
                try:
                    os.remove(temp_file)
                except:
                    pass
            # Return current state, but update profanity info
            stream_results["profanity_info"] = "No speech detected. Keep talking..."
            return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
        
        # Update transcript
        stream_results["transcript"] = transcript
        
        # Analyze for profanity
        analysis = detect_profanity(transcript, threshold=0.5)
        
        # Check if profanity was detected
        if analysis.get("profanity", False):
            profane_words = ", ".join(analysis.get("profane_words", []))
            stream_results["profanity_info"] = f"Profanity Detected (Score: {analysis['score']:.2f})\nProfane Words: {profane_words}"
            
            # Rephrase to clean text
            clean_text = rephrase_profanity(transcript)
            stream_results["clean_text"] = clean_text
            
            # Create audio from cleaned text
            audio_file = text_to_speech(clean_text)
            if audio_file:
                stream_results["audio_output"] = audio_file
        else:
            stream_results["profanity_info"] = f"No Profanity Detected (Score: {analysis['score']:.2f})"
            stream_results["clean_text"] = transcript
            
            # Use original text for audio if no profanity
            audio_file = text_to_speech(transcript)
            if audio_file:
                stream_results["audio_output"] = audio_file
        
        # Clean up temporary file if we created it
        if temp_file and temp_file.startswith("temp_stream_") and os.path.exists(temp_file):
            try:
                os.remove(temp_file)
            except:
                pass
        
        return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]
        
    except Exception as e:
        error_msg = f"Error processing streaming audio: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        
        # Make sure all models are on CPU if in ZeroGPU mode
        if IS_ZEROGPU and torch.cuda.is_available():
            try:
                whisper_model.to(torch.device("cpu"))
                profanity_model.to(torch.device("cpu"))
                t5_model.to(torch.device("cpu"))
                tts_model.to(torch.device("cpu"))
                vocoder.to(torch.device("cpu"))
            except:
                pass
        
        # Update profanity info with error message
        stream_results["profanity_info"] = f"Error: {str(e)}"
        
        return stream_results["transcript"], stream_results["profanity_info"], stream_results["clean_text"], stream_results["audio_output"]

def start_streaming():
    """Start the real-time audio processing"""
    global processing_active, stream_results
    
    if not models_loaded:
        return "Models not loaded yet. Please wait for initialization to complete."
    
    if processing_active:
        return "Streaming is already active."
    
    # Reset results
    stream_results = {
        "transcript": "",
        "profanity_info": "Waiting for audio input...",
        "clean_text": "",
        "audio_output": None
    }
    
    processing_active = True
    logger.info("Started real-time audio processing")
    return "Started real-time audio processing. Speak into your microphone."

def stop_streaming():
    """Stop the real-time audio processing"""
    global processing_active
    
    if not processing_active:
        return "Streaming is not active."
    
    processing_active = False
    return "Stopped real-time audio processing."

def create_ui():
    """Create the Gradio UI"""
    # Simple CSS for styling
    css = """
    /* Fix for dark mode text visibility */
    .dark .gr-input, 
    .dark textarea,
    .dark .gr-textbox,
    .dark [data-testid="textbox"] {
        color: white !important;
        background-color: #2c303b !important;
    }

    .dark .gr-box, 
    .dark .gr-form, 
    .dark .gr-panel, 
    .dark .gr-block {
        color: white !important;
    }
    
    /* Highlighted text container - with dark mode fixes */
    .highlighted-text {
        border: 1px solid #ddd;
        border-radius: 5px;
        padding: 10px;
        margin: 10px 0;
        background-color: #f9f9f9;
        font-family: sans-serif;
        max-height: 300px;
        overflow-y: auto;
        color: #333 !important; /* Ensure text is dark for light mode */
    }
    
    /* Dark mode specific styling for highlighted text */
    .dark .highlighted-text {
        background-color: #2c303b !important;
        color: #ffffff !important;
        border-color: #4a4f5a !important;
    }
    
    /* Make sure text in the highlighted container remains visible in both themes */
    .highlighted-text, .dark .highlighted-text {
        color-scheme: light dark;
    }
    
    /* Loading animation */
    .loading {
        display: inline-block;
        width: 20px;
        height: 20px;
        border: 3px solid rgba(0,0,0,.3);
        border-radius: 50%;
        border-top-color: #3498db;
        animation: spin 1s ease-in-out infinite;
    }
    
    @keyframes spin {
        to { transform: rotate(360deg); }
    }
    """

    # Create a custom theme based on Soft but explicitly set to light mode
    light_theme = gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="blue",
        neutral_hue="gray"
    )

    # Set theme to light mode and disable theme switching
    with gr.Blocks(css=css, theme=light_theme, analytics_enabled=False) as ui:
        # Model initialization
        init_status = gr.State("")
        
        gr.Markdown(
            """
            # Profanity Detection & Replacement System
            Detect, rephrase, and listen to cleaned content from text or audio!
            """,
            elem_classes="header"
        )
        
        # The rest of your UI code remains unchanged...
        # Initialize models button with status indicators
        with gr.Row():
            with gr.Column(scale=3):
                init_button = gr.Button("Initialize Models", variant="primary")
                init_output = gr.Textbox(label="Initialization Status", interactive=False)
            with gr.Column(scale=1):
                model_status = gr.HTML(
                    """<div style="text-align: center; padding: 5px;">
                    <p><b>Model Status:</b> <span style="color: #e74c3c;">Not Loaded</span></p>
                    </div>"""
                )
        
        # Global sensitivity slider
        sensitivity = gr.Slider(
            minimum=0.2,
            maximum=0.95,
            value=0.5,
            step=0.05,
            label="Profanity Detection Sensitivity",
            info="Lower values are more permissive, higher values are more strict"
        )

        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("### Choose an Input Method")

        # Text Analysis
        with gr.Tabs():
            with gr.TabItem("Text Analysis", elem_id="text-tab"):
                with gr.Row():
                    text_input = gr.Textbox(
                        label="Enter Text",
                        placeholder="Type your text here...",
                        lines=5,
                        elem_classes="textbox"
                    )
                with gr.Row():
                    text_button = gr.Button("Analyze Text", variant="primary")
                    clear_button = gr.Button("Clear", variant="secondary")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        text_output = gr.Textbox(label="Results", lines=10)
                        highlighted_output = gr.HTML(label="Detected Profanity", elem_classes="highlighted-text")
                    with gr.Column(scale=1):
                        text_audio_output = gr.Audio(label="Rephrased Audio", type="filepath")

            # Audio Analysis
            with gr.TabItem("Audio Analysis", elem_id="audio-tab"):
                gr.Markdown("### Upload or Record Audio")
                audio_input = gr.Audio(
                    label="Audio Input",
                    type="filepath",
                    sources=["microphone", "upload"]
                    #waveform_options=gr.WaveformOptions(waveform_color="#4a90e2")
                )
                with gr.Row():
                    audio_button = gr.Button("Analyze Audio", variant="primary")
                    clear_audio_button = gr.Button("Clear", variant="secondary")
                
                with gr.Row():
                    with gr.Column(scale=2):
                        audio_output = gr.Textbox(label="Results", lines=10, show_copy_button=True)
                        audio_highlighted_output = gr.HTML(label="Detected Profanity", elem_classes="highlighted-text")
                    with gr.Column(scale=1):
                        clean_audio_output = gr.Audio(label="Rephrased Audio", type="filepath")
            
            # Real-time Streaming
            with gr.TabItem("Real-time Streaming", elem_id="streaming-tab"):
                gr.Markdown("### Real-time Audio Processing")
                gr.Markdown("Enable real-time audio processing to filter profanity as you speak.")
                
                with gr.Row():
                    with gr.Column(scale=1):
                        start_stream_button = gr.Button("Start Real-time Processing", variant="primary")
                        stop_stream_button = gr.Button("Stop Real-time Processing", variant="secondary")
                        stream_status = gr.Textbox(label="Streaming Status", value="Inactive", interactive=False)
                        
                        # Add microphone input specifically for streaming
                        stream_audio_input = gr.Audio(
                            label="Streaming Microphone Input",
                            type="filepath",
                            sources=["microphone"],
                            streaming=True
                            #waveform_options=gr.WaveformOptions(waveform_color="#4a90e2")
                        )
                    
                    with gr.Column(scale=2):
                        # Add elements to display streaming results
                        stream_transcript = gr.Textbox(label="Live Transcription", lines=2)
                        stream_profanity_info = gr.Textbox(label="Profanity Detection", lines=2)
                        stream_clean_text = gr.Textbox(label="Clean Text", lines=2)
                        # Element to play the clean audio
                        stream_audio_output = gr.Audio(label="Clean Audio Output", type="filepath")
                
                gr.Markdown("""
                ### How Real-time Streaming Works
                1. Click "Start Real-time Processing" to begin
                2. Use the microphone input to speak
                3. The system will process audio in real-time, detect and clean profanity
                4. You'll see the transcription, profanity info, and clean output appear above
                5. Click "Stop Real-time Processing" when finished
                
                Note: This feature requires microphone access and may have some latency.
                """)

        # Event handlers
        def update_model_status(status_text):
            """Update both the status text and the visual indicator"""
            if "successfully" in status_text.lower():
                status_html = """<div style="text-align: center; padding: 5px;">
                <p><b>Model Status:</b> <span style="color: #2ecc71;">Loaded ✓</span></p>
                </div>"""
            elif "error" in status_text.lower():
                status_html = """<div style="text-align: center; padding: 5px;">
                <p><b>Model Status:</b> <span style="color: #e74c3c;">Error ✗</span></p>
                </div>"""
            else:
                status_html = """<div style="text-align: center; padding: 5px;">
                <p><b>Model Status:</b> <span style="color: #f39c12;">Loading...</span></p>
                </div>"""
            return status_text, status_html
            
        init_button.click(
            lambda: update_model_status("Loading models, please wait..."),
            inputs=[],
            outputs=[init_output, model_status]
        ).then(
            load_models,
            inputs=[],
            outputs=[init_output]
        ).then(
            update_model_status,
            inputs=[init_output],
            outputs=[init_output, model_status]
        )
        
        text_button.click(
            text_analysis, 
            inputs=[text_input, sensitivity], 
            outputs=[text_output, highlighted_output, text_audio_output]
        )
        
        clear_button.click(
            lambda: [None, None, None], 
            inputs=None, 
            outputs=[text_input, highlighted_output, text_audio_output]
        )

        audio_button.click(
            analyze_audio, 
            inputs=[audio_input, sensitivity], 
            outputs=[audio_output, audio_highlighted_output, clean_audio_output]
        )

        clear_audio_button.click(
            lambda: [None, None, None, None], 
            inputs=None, 
            outputs=[audio_input, audio_output, audio_highlighted_output, clean_audio_output]
        )
        
        start_stream_button.click(
            start_streaming,
            inputs=[],
            outputs=[stream_status]
        )
        
        stop_stream_button.click(
            stop_streaming,
            inputs=[],
            outputs=[stream_status]
        )
        
        # Connect the streaming audio input to our processing function
        # First function to debug the audio chunk format
        def debug_audio_format(audio_chunk):
            """Debug function to log audio format"""
            format_info = f"Type: {type(audio_chunk)}"
            if isinstance(audio_chunk, tuple):
                format_info += f", Length: {len(audio_chunk)}"
                for i, item in enumerate(audio_chunk):
                    format_info += f", Item {i} type: {type(item)}"
            logger.info(f"Audio chunk format: {format_info}")
            return audio_chunk
            
        # Use the stream method with preprocessor for debugging
        stream_audio_input.stream(
            fn=process_stream_chunk,
            inputs=[stream_audio_input], 
            outputs=[stream_transcript, stream_profanity_info, stream_clean_text, stream_audio_output],
            preprocess=debug_audio_format
        )
        
    return ui

if __name__ == "__main__":
    # Set environment variable to avoid OpenMP conflicts
    os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
    
    # Create and launch the UI
    ui = create_ui()
    ui.launch(server_name="0.0.0.0", share=True)