Sync SafeSeal app

README.md

@@ -1,13 +1,72 @@
 ---
-title: SafeSeal
-emoji: 🏆
-colorFrom: purple
+title: SafeSeal Watermark
+emoji: 🔒
+colorFrom: blue
 colorTo: purple
-sdk: gradio
-sdk_version: 5.49.1
+sdk: streamlit
+sdk_version: 1.50.0
 app_file: app.py
 pinned: false
-short_description: Demo for SafeSeal watermarking method.
+license: mit
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# SafeSeal Watermark
+
+**Content-Preserving Watermarking for Large Language Model Deployments.**
+
+Generate watermarked text by key-conditioned sampling words with context-aware synonyms.
+
+## Features
+
+- 🔑 **Secret Key**: Deterministic watermarking with user-controlled key
+- 📊 **BERTScore Filtering**: Adjustable similarity threshold (0.0 - 1.0)
+- 🏆 **Tournament Sampling**: Select synonyms using tournament-based randomization
+- ✨ **Visual Highlighting**: See exactly which words were changed
+- 🚀 **GPU Support**: Fast inference with automatic GPU detection
+- 🛡️ **Smart Filtering**: Excludes antonyms, specific nouns in same category, and preserves entity names
+
+## How It Works
+
+1. **Entity Detection**: Extracts eligible words (nouns, verbs, adjectives, adverbs) while skipping named entities
+2. **Candidate Generation**: Uses RoBERTa-base to generate semantically similar alternatives
+3. **BERTScore Filtering**: Evaluates candidates against a similarity threshold
+4. **Tournament Selection**: Deterministically selects replacements based on secret key
+5. **Visualization**: Highlights changed words in the output
+
+## Usage
+
+1. Enter your text in the left panel
+2. Adjust hyperparameters in the sidebar:
+   - **Secret Key**: Used for deterministic randomization
+   - **Threshold**: Similarity threshold (default: 0.98)
+   - **Tournament parameters**: Fine-tune the selection process
+3. Click "🚀 Generate Watermark"
+4. View the watermarked text with highlighted changes
+
+## Parameters
+
+- **Secret Key**: Used for deterministic randomization
+- **Threshold (0.98)**: BERTScore similarity threshold - higher = more conservative changes
+- **m (10)**: Number of tournament rounds
+- **c (2)**: Competitors per tournament match
+- **h (6)**: Context size (left tokens to consider)
+- **Alpha (1.1)**: Temperature scaling factor
+
+## Technical Details
+
+- **Model**: RoBERTa-base for masked language modeling
+- **Similarity Scoring**: BERTScore F1 scores
+- **Selection**: Tournament-based deterministic sampling
+- **Filtering**: POS tag matching, antonym exclusion, semantic compatibility checks
+
+## Example
+
+**Input:**
+> "The quick brown fox jumps over the lazy dog."
+
+**Watermarked Output:**
+> "The swift brown fox leaps over the idle dog."
+
+Changed words highlighted: swift (was quick), leaps (was jumps), idle (was lazy)
+
+⚠️ **Demo Version**: This is a demonstration using a light model to showcase the watermarking pipeline. Results may not be perfect and are intended for testing purposes only.

SynthID_randomization.py

@@ -0,0 +1,294 @@
+# Sampling with Tournament
+# input:
+#   original: original word
+#   alternatives: substitute candidates
+#   similarity: similarity scores 
+#   context: 1-token left context
+#   key: secret key for hashing
+#   m: number of tournament rounds. Set to 2
+#   c: number of candidates per round (not including original): c is set to the number of alternatives = K
+#   alpha: scaling factor like temperature. Set to 1
+
+#   Process:
+# 1. Normalize similarity scores to [0,1] though softmax with scaling factor alpha
+# 2. Create random score for each candidate with m rounds: so we have a matrix of shape (K, m)
+# 3. Each score is generated by hashing key, context, candidate, round index
+# 4. For the first round, sampling from normalized similarity scores M = K^m candidates (each candidate can appear multiple times) and divide into groups of size K
+# 5. For each group, select the candidate with the highest random score to go to the next round. If they are tied, just randomly select one of the tied candidates. For example, K = 3, we have 1 1 0 score, so break the 1 1 tie randomly.
+# 6: For next round, use the winners from the previous round, but the scores from the matrix for that round. Repeat until we have one winner.
+
+# Output: return winner candidate index
+
+
+import hmac
+import hashlib
+import math
+from typing import List, Tuple
+
+def _softmax(xs: List[float], alpha: float = 1.0) -> List[float]:
+    # Temperature-scaled softmax normalized to sum 1
+    if not xs:
+        return []
+    # subtract max for numerical stability
+    m = max(xs)
+    zs = [math.exp(alpha * (x - m)) for x in xs]
+    s = sum(zs)
+    return [z / s for z in zs]
+
+def _hash_bytes(key: str, payload: str) -> bytes:
+    return hmac.new(key.encode("utf-8"), payload.encode("utf-8"), hashlib.sha256).digest()
+
+def _hash_to_uniform01(key: str, payload: str) -> float:
+    # Map 8 bytes to [0,1) as a 64-bit integer / 2**64
+    b = _hash_bytes(key, payload)[:8]
+    n = int.from_bytes(b, "big", signed=False)
+    return n / 2**64
+
+def _ctx_str(context: List[str]) -> str:
+    # Use last token (you can widen to last-4 if you like)
+    toks = (context or [])
+    return " ".join(toks[-1:]).lower()
+
+def _per_round_score(key: str, ctx: str, candidate: str, round_idx: int) -> int:
+    # Deterministic "random" score following Bernoulli(0.5): 0 or 1
+    payload = f"score::{ctx}::{candidate}::r{round_idx}"
+    uniform_score = _hash_to_uniform01(key, payload)
+    # Convert uniform [0,1) to Bernoulli(0.5): 0 if < 0.5, 1 if >= 0.5
+    return 1 if uniform_score >= 0.5 else 0
+
+def _sample_categorical_from_uniform(u: float, probs: List[float]) -> int:
+    # Inverse-CDF sampling from probs using a uniform u in [0,1)
+    c = 0.0
+    for i, p in enumerate(probs):
+        c += p
+        if u < c:
+            return i
+    return len(probs) - 1  # fallback on last due to float sums
+
+def _draw_M_candidates(key: str, ctx: str, probs: List[float], K: int, c: int, m: int) -> List[int]:
+    """
+    Draw M = c^m candidates (indices 0..K-1) with replacement from probs,
+    using a deterministic stream of uniforms keyed by (key, ctx).
+    """
+    M = c ** m
+    picks = []
+    for draw_idx in range(M):
+        u = _hash_to_uniform01(key, f"draw::{ctx}::m{m}::c{c}::i{draw_idx}")
+        j = _sample_categorical_from_uniform(u, probs)
+        picks.append(j)
+    return picks
+
+def _run_tournament_round(
+    key: str,
+    ctx: str,
+    picks: List[int],
+    candidates: List[str],
+    round_idx: int,
+    group_size: int,
+) -> List[int]:
+    """
+    Split 'picks' into consecutive groups of size 'group_size'.
+    For each group, select the winner = argmax per-round-score;
+    tie-break deterministically.
+    Returns the list of winning indices (into candidates list).
+    """
+    assert len(picks) % group_size == 0, "Group partition must divide evenly."
+    winners = []
+    
+    # Add detailed logging for all rounds
+    # print(f"\n=== ROUND {round_idx} DETAILS ===")
+    # print(f"Total picks: {len(picks)}")
+    # print(f"Group size: {group_size}")
+    # print(f"Number of groups: {len(picks) // group_size}")
+    # print(f"Picks array: {picks}")
+    # print(f"Candidates: {candidates}")
+    
+    for g in range(0, len(picks), group_size):
+        group = picks[g:g+group_size]  # indices into candidates list
+        group_num = g // group_size + 1
+        
+        # Compute per-round scores
+        scored: List[Tuple[int, int]] = []
+        for idx in group:
+            cand = candidates[idx]
+            s = _per_round_score(key, ctx, cand, round_idx)
+            scored.append((s, idx))
+        
+        # Add detailed logging for all rounds
+        # print(f"\n--- Group {group_num} ---")
+        # print(f"Group indices: {group}")
+        # print(f"Group candidates: {[candidates[idx] for idx in group]}")
+        # print("Scores:")
+        # for score, idx in scored:
+        #     print(f"  {candidates[idx]}: {score}")
+        
+        # Find max score; tie-break with a deterministic secondary key
+        max_s = max(s for s, _ in scored)
+        tied = [idx for (s, idx) in scored if s == max_s]
+        
+        if len(tied) == 1:
+            winner_idx = tied[0]
+            winners.append(winner_idx)
+            # print(f"Winner: {candidates[winner_idx]} (score: {max_s})")
+        else:
+            # Random tie-breaker: pick one randomly from tied candidates
+            import random
+            # Create a more robust seed by combining key, context, round, and group info
+            seed_string = f"{key}:{ctx}:r{round_idx}:g{group_num}"
+            # Convert string to integer seed using a more reliable method
+            seed_value = sum(ord(c) * (i + 1) for i, c in enumerate(seed_string))
+            random.seed(seed_value)
+            winner_idx = random.choice(tied)
+            random.seed()  # Reset seed to avoid affecting other random operations
+            winners.append(winner_idx)
+            # print(f"TIE! Winners: {[candidates[idx] for idx in tied]}")
+            # print(f"Random tie-breaker winner: {candidates[winner_idx]} (score: {max_s})")
+            # print(f"  (Seed used: {seed_string} -> {seed_value})")
+    
+    # Add summary logging for all rounds
+    # print(f"\nRound {round_idx} winners: {[candidates[idx] for idx in winners]}")
+    # print(f"Winner indices: {winners}")
+    # print("=" * 50)
+    
+    return winners
+
+def _show_score_matrix(key: str, ctx: str, candidates: List[str], m: int):
+    """
+    Display the complete score matrix for all candidates across all rounds.
+    Each candidate gets a Bernoulli(0.5) score for each round.
+    """
+    # print(f"\n=== SCORE MATRIX (Bernoulli 0.5) ===")
+    # print(f"Format: candidate -> [round1_score, round2_score, ...]")
+    # for i, candidate in enumerate(candidates):
+    #     scores = []
+    #     for round_idx in range(1, m + 1):
+    #         score = _per_round_score(key, ctx, candidate, round_idx)
+    #         scores.append(score)
+    #     print(f"{candidate:8} -> {scores}")
+    # print("=" * 50)
+    pass
+
+def tournament_randomize(
+    original: str,
+    alternatives: List[str],
+    similarity: List[float],
+    context: List[str],
+    key: str,
+    m: int = 2,
+    c: int = 2,
+    alpha: float = 1.0,
+) -> int:
+    """
+    Tournament sampling among alternatives only (K = len(alternatives)).
+    Returns the WINNER INDEX (0..K-1) into 'alternatives'.
+
+    Steps:
+      1) softmax(similarity, alpha) -> probs over alternatives
+      2) build per-round scores for each candidate via HMAC (on the fly)
+      3) Round 1: draw M = c^m picks from probs; group into size c; pick per-group winners by scores of round 1
+      4) Next rounds: group winners into size c; pick winners using that round's scores
+      5) Repeat until one winner remains
+    """
+    assert len(alternatives) == len(similarity) > 0, "Need at least one alternative with a similarity score."
+    assert m >= 1, "Tournament rounds m must be >= 1"
+    assert c >= 2, "Number of competitors per match c must be >= 2"
+
+    K = len(alternatives)
+    ctx = _ctx_str(context)
+    probs = _softmax(similarity, alpha=alpha)
+
+    # Add detailed logging for setup
+    # print(f"\n=== TOURNAMENT SETUP ===")
+    # print(f"Original word: '{original}'")
+    # print(f"Alternatives: {alternatives}")
+    # print(f"Similarity scores: {similarity}")
+    # print(f"Context: {context}")
+    # print(f"Key: {key}")
+    # print(f"Tournament rounds (m): {m}")
+    # print(f"Competitors per match (c): {c}")
+    # print(f"Alpha (temperature): {alpha}")
+    # print(f"K (number of alternatives): {K}")
+    # print(f"Context string: '{ctx}'")
+    # print(f"Normalized probabilities: {[f'{p:.6f}' for p in probs]}")
+    # print(f"Expected picks (M = c^m): {c**m}")
+
+    # First round: M = c^m draws from probs
+    picks = _draw_M_candidates(key, ctx, probs, K, c, m)  # indices 0..K-1 into alternatives
+    
+    # Show the picks array for the first round
+    # print(f"Generated picks array: {picks}")
+    # print(f"Picks correspond to words: {[alternatives[idx] for idx in picks]}")
+
+    # We treat 'candidates' as the alternatives list; indices map directly
+    candidates = alternatives
+    
+    # Show the complete score matrix for all candidates across all rounds
+    _show_score_matrix(key, ctx, candidates, m)
+
+    # Run m rounds. Each round groups current list into blocks of size c,
+    # picks one per group using per-round scores.
+    current = picks
+    for r in range(1, m + 1):
+        # group_size is now c instead of K
+        current = _run_tournament_round(
+            key=key,
+            ctx=ctx,
+            picks=current,
+            candidates=candidates,
+            round_idx=r,
+            group_size=c,
+        )
+        # After each round, the number of survivors shrinks by factor c.
+        # After m rounds, we must have exactly one winner.
+    assert len(current) == 1, f"Expected a single winner after {m} rounds, got {len(current)}"
+    
+    final_winner_idx = current[0]
+    # print(f"\n=== FINAL RESULT ===")
+    # print(f"Winner index: {final_winner_idx}")
+    # print(f"Winner word: '{candidates[final_winner_idx]}'")
+    
+    return final_winner_idx
+
+# Convenience wrapper that returns the selected word
+def tournament_select_word(
+    original: str,
+    alternatives: List[str],
+    similarity: List[float],
+    context: List[str],
+    key: str,
+    m: int = 2,
+    c: int = 2,
+    alpha: float = 1.0,
+) -> str:
+    idx = tournament_randomize(original, alternatives, similarity, context, key, m, c, alpha)
+    return alternatives[idx]
+
+
+# # Sample running
+# original = "big"
+# alts = ["large", "huge", "massive"]          # K = 3
+# sims  = [0.82, 0.55, 0.60]                   # any similarity scores
+# ctx   = ["a"]                                 # 1-token left context (e.g., "... a ___")
+# key   = "super-secret-key"
+# m     = 2
+# c     = 2
+# alpha = 1.0
+
+# winner_idx = tournament_randomize(original, alts, sims, ctx, key, m, c, alpha)
+# print("Winner index:", winner_idx, "->", alts[winner_idx])
+
+# # Or directly:
+# print("Winner word:", tournament_select_word(original, alts, sims, ctx, key, m, c, alpha))
+
+# # Test with a different key
+# key2 = "different-secret-key"
+# winner_idx2 = tournament_randomize(original, alts, sims, ctx, key2, m, c, alpha)
+# print("Winner index with different key:", winner_idx2, "->", alts[winner_idx2])
+
+# # Test with a different key
+# key3 = "key224242"
+# winner_idx3 = tournament_randomize(original, alts, sims, ctx, key3, m, c, alpha)
+# print("Winner index with different key:", winner_idx3, "->", alts[winner_idx3])
+
+
+

app.py

@@ -0,0 +1,364 @@
+import streamlit as st
+import os
+import torch
+from transformers import RobertaTokenizer, RobertaForMaskedLM
+import spacy
+import subprocess
+import sys
+import nltk
+from nltk.tokenize import word_tokenize
+from utils_final import extract_entities_and_pos, whole_context_process_sentence
+
+# Download NLTK data if not available
+def setup_nltk():
+    """Setup NLTK data with error handling."""
+    try:
+        nltk.download('punkt_tab', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('averaged_perceptron_tagger_eng', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('wordnet', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('omw-1.4', quiet=True)
+    except:
+        pass
+
+setup_nltk()
+
+# Set environment
+cache_dir = '/network/rit/lab/Lai_ReSecureAI/kiel/wmm'
+
+# Load spaCy model - download if not available
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    print("Downloading spaCy model...")
+    subprocess.check_call([sys.executable, "-m", "spacy", "download", "en_core_web_sm"])
+    nlp = spacy.load("en_core_web_sm")
+
+# Define apply_replacements function (from Safeseal_gen_final.py)
+def apply_replacements(sentence, replacements):
+    """
+    Apply replacements to the sentence while preserving original formatting, spacing, and punctuation.
+    """
+    doc = nlp(sentence)  # Tokenize the sentence
+    tokens = [token.text_with_ws for token in doc]  # Preserve original whitespace with tokens
+
+    # Apply replacements based on token positions
+    for position, target, replacement in replacements:
+        if position < len(tokens) and tokens[position].strip() == target:
+            tokens[position] = replacement + (" " if tokens[position].endswith(" ") else "")
+
+    # Reassemble the sentence
+    return "".join(tokens)
+
+# Initialize session state for model caching
+@st.cache_resource
+def load_model():
+    """Load the model and tokenizer (cached to avoid reloading on every run)"""
+    print("Loading model...")
+    tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+    lm_model = RobertaForMaskedLM.from_pretrained('roberta-base', attn_implementation="eager")
+    
+    tokenizer.model_max_length = 512
+    tokenizer.max_len = 512
+    
+    if hasattr(lm_model.config, 'max_position_embeddings'):
+        lm_model.config.max_position_embeddings = 512
+    
+    lm_model.eval()
+    
+    if torch.cuda.is_available():
+        lm_model = lm_model.cuda()
+        print(f"Model loaded on GPU: {torch.cuda.get_device_name()}")
+    else:
+        print("Model loaded on CPU")
+    
+    return tokenizer, lm_model
+
+sampling_results = []
+
+def process_text_wrapper(text, tokenizer, lm_model, Top_K, threshold, secret_key, m, c, h, alpha, batch_size=32, max_length=512, similarity_context_mode='whole'):
+    """
+    Wrapper function to process text and return watermarked output with tracking of changes.
+    """
+    global sampling_results
+    sampling_results = []
+    
+    lines = text.splitlines(keepends=True)
+    final_text = []
+    total_randomized_words = 0
+    total_words = len(word_tokenize(text))
+    
+    # Track changed words and their positions
+    changed_words = []  # List of (original, replacement, position)
+    
+    for line in lines:
+        if line.strip():
+            replacements, sampling_results_line = whole_context_process_sentence(
+                text,
+                line.strip(),
+                tokenizer, lm_model, Top_K, threshold,
+                secret_key, m, c, h, alpha, "output",
+                batch_size=batch_size, max_length=max_length, similarity_context_mode=similarity_context_mode
+            )
+            
+            sampling_results.extend(sampling_results_line)
+            
+            if replacements:
+                randomized_line = apply_replacements(line, replacements)
+                final_text.append(randomized_line)
+                
+                # Track ONLY actual changes (where original != replacement)
+                for position, original, replacement in replacements:
+                    if original != replacement:
+                        changed_words.append((original, replacement, position))
+                        total_randomized_words += 1
+            else:
+                final_text.append(line)
+        else:
+            final_text.append(line)
+    
+    return "".join(final_text), total_randomized_words, total_words, changed_words, sampling_results
+
+def create_html_with_highlights(original_text, watermarked_text, changed_words_info, sampling_results):
+    """
+    Create HTML with highlighted changed words using spaCy tokenization.
+    """
+    # Create a set of replacement words that were actually changed (not same as original)
+    actual_replacements = set()
+    replacement_to_original = {}
+    
+    for original, replacement, _ in changed_words_info:
+        if original.lower() != replacement.lower():  # Only map actual changes
+            actual_replacements.add(replacement.lower())
+            replacement_to_original[replacement.lower()] = original
+    
+    # Parse watermarked text with spaCy
+    doc_watermarked = nlp(watermarked_text)
+    
+    # Build HTML by processing the watermarked text
+    result_html = []
+    words_highlighted = set()  # Track which words we've highlighted (to avoid duplicates)
+    
+    for token in doc_watermarked:
+        text = token.text_with_ws
+        text_clean = token.text.strip('.,!?;:')
+        text_lower = text_clean.lower()
+        
+        # Only highlight if this word is in our actual replacements set
+        # and we haven't already highlighted this exact word
+        if text_lower in actual_replacements and text_lower not in words_highlighted:
+            original_word = replacement_to_original.get(text_lower, text_clean)
+            
+            # Only highlight if actually different from original
+            if original_word.lower() != text_lower:
+                tooltip = f"Original: {original_word} → New: {text_clean}"
+                # Enhanced highlighting with better colors
+                highlighted_text = f"<mark style='background: linear-gradient(120deg, #84fab0 0%, #8fd3f4 100%); padding: 2px 6px; border-radius: 4px; font-weight: 500; box-shadow: 0 1px 2px rgba(0,0,0,0.1);' title='{tooltip}'>{text_clean}</mark>"
+                
+                # Preserve trailing whitespace and punctuation
+                if text != text_clean:
+                    highlighted_text += text[len(text_clean):]
+                
+                result_html.append(highlighted_text)
+                words_highlighted.add(text_lower)  # Mark as highlighted
+            else:
+                result_html.append(text)
+        else:
+            result_html.append(text)
+    
+    # Return just the inner content without the outer div (added by caller)
+    return "".join(result_html)
+
+# Streamlit UI
+def main():
+    st.set_page_config(
+        page_title="Watermarked Text Generator",
+        page_icon="🔒",
+        layout="wide"
+    )
+    
+    # Centered and styled title
+    st.markdown(
+        """
+        <div style="text-align: center; margin-bottom: 10px;">
+            <h1 style="color: #4A90E2; font-size: 2.5rem; font-weight: bold; margin: 0;">
+                🔒 SafeSeal Watermark
+            </h1>
+        </div>
+        <div style="text-align: center; margin-bottom: 20px; color: #666; font-size: 1.1rem;">
+            Content-Preserving Watermarking for Large Language Model Deployments.
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+    
+    # Add a nice separator
+    st.markdown("---")
+    
+    # Sidebar for hyperparameters
+    with st.sidebar:
+        st.markdown("### ⚙️ Hyperparameters")
+        st.caption("Configure the watermarking algorithm")
+        
+        # Main inputs
+        secret_key = st.text_input(
+            "🔑 Secret Key",
+            value="My_Secret_Key",
+            help="Secret key for deterministic randomization"
+        )
+        
+        threshold = st.slider(
+            "📊 Similarity Threshold",
+            min_value=0.0,
+            max_value=1.0,
+            value=0.98,
+            step=0.01,
+            help="BERTScore similarity threshold (higher = more similar replacements)"
+        )
+        
+        st.divider()
+        
+        # Tournament Sampling parameters
+        st.markdown("### 🏆 Tournament Sampling")
+        st.caption("Control the randomization process")
+        
+        # Hidden Top_K parameter (default 6)
+        Top_K = 6
+        
+        m = st.number_input(
+            "m (Tournament Rounds)",
+            min_value=1,
+            max_value=20,
+            value=10,
+            help="Number of tournament rounds"
+        )
+        
+        c = st.number_input(
+            "c (Competitors per Round)",
+            min_value=2,
+            max_value=10,
+            value=2,
+            help="Number of competitors per tournament match"
+        )
+        
+        h = st.number_input(
+            "h (Context Size)",
+            min_value=1,
+            max_value=20,
+            value=6,
+            help="Number of left context tokens to consider"
+        )
+        
+        alpha = st.slider(
+            "Alpha (Temperature)",
+            min_value=0.1,
+            max_value=5.0,
+            value=1.1,
+            step=0.1,
+            help="Temperature scaling factor for softmax"
+        )
+    
+    # Main content area
+    col1, col2 = st.columns(2)
+    
+    # Check if model is loaded
+    if 'tokenizer' not in st.session_state:
+        with st.spinner("Loading model... This may take a minute"):
+            tokenizer, lm_model = load_model()
+            st.session_state.tokenizer = tokenizer
+            st.session_state.lm_model = lm_model
+    
+    with col1:
+        st.markdown("### 📝 Input Text")
+        input_text = st.text_area(
+            "Enter text to watermark",
+            height=400,
+            placeholder="Paste your text here to generate a watermarked version...",
+            label_visibility="collapsed"
+        )
+        
+        # Process button at the bottom of input column
+        if st.button("🚀 Generate Watermark", type="primary", use_container_width=True):
+            if not input_text or len(input_text.strip()) == 0:
+                st.warning("Please enter some text to watermark.")
+            else:
+                with st.spinner("Generating watermarked text... This may take a few moments"):
+                    try:
+                        # Process the text
+                        watermarked_text, total_randomized_words, total_words, changed_words, sampling_results = process_text_wrapper(
+                            input_text,
+                            st.session_state.tokenizer,
+                            st.session_state.lm_model,
+                            Top_K=int(Top_K),
+                            threshold=float(threshold),
+                            secret_key=secret_key,
+                            m=int(m),
+                            c=int(c),
+                            h=int(h),
+                            alpha=float(alpha),
+                            batch_size=32,
+                            max_length=512,
+                            similarity_context_mode='whole'
+                        )
+                        
+                        # Store results in session state
+                        st.session_state.watermarked_text = watermarked_text
+                        st.session_state.changed_words = changed_words
+                        st.session_state.sampling_results = sampling_results
+                        st.session_state.total_randomized = total_randomized_words
+                        st.session_state.total_words = total_words
+                        
+                        st.success(f"Watermark generated! Changed {total_randomized_words} out of {total_words} words ({100*total_randomized_words/max(total_words,1):.1f}%)")
+                    except Exception as e:
+                        st.error(f"Error generating watermark: {str(e)}")
+                        import traceback
+                        st.code(traceback.format_exc())
+    
+    with col2:
+        st.markdown("### 🔒 Watermarked Text")
+        
+        # Display watermarked text with highlights
+        if 'watermarked_text' in st.session_state:
+            highlight_html = create_html_with_highlights(
+                input_text,
+                st.session_state.watermarked_text,
+                st.session_state.changed_words,
+                st.session_state.sampling_results
+            )
+            # Show highlighted version with border - wrap the complete HTML
+            full_html = f"""
+            <div style='padding: 15px; background-color: #f8f9fa; border-radius: 8px; border: 1px solid #e0e0e0; min-height: 400px; max-height: 400px; overflow-y: auto; line-height: 1.8; font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif; font-size: 15px; white-space: pre-wrap; word-wrap: break-word;'>
+            {highlight_html}
+            </div>
+            """
+            st.markdown(full_html, unsafe_allow_html=True)
+        else:
+            st.info("👈 Enter text in the left panel and click 'Generate Watermark' to start")
+    
+    # Footer
+    st.divider()
+    st.caption("🔒 Secure AI Watermarking Tool | Built with SafeSeal")
+    
+    # Demo warning at the bottom
+    st.markdown(
+        """
+        <div style="text-align: center; margin-top: 20px; padding: 10px; font-size: 0.85rem; color: #666;">
+            ⚠️ <strong>Demo Version</strong>: This is a demonstration using a light model to showcase the watermarking pipeline. 
+            Results may not be perfect and are intended for testing purposes only.
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+if __name__ == "__main__":
+    # Run the app
+    main()
+

requirements.txt

@@ -0,0 +1,10 @@
+streamlit>=1.28.0
+torch>=2.0.0
+transformers>=4.30.0
+spacy>=3.4.0
+nltk>=3.8.0
+bert-score>=0.3.13
+pandas>=2.0.0
+numpy>=1.24.0
+scikit-learn>=1.3.0
+

utils_final.py

@@ -0,0 +1,1213 @@
+import os
+import re
+import time
+import math
+import torch
+import string
+import spacy
+import pandas as pd
+import numpy as np
+import nltk
+import sys
+import subprocess
+from nltk.tokenize import word_tokenize
+from nltk.stem.wordnet import WordNetLemmatizer
+from nltk.corpus import wordnet as wn
+import json
+from filelock import FileLock
+from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor, as_completed
+from functools import lru_cache
+from typing import List, Tuple, Dict, Any
+import multiprocessing as mp
+
+# Ensure the HF_HOME environment variable points to your desired cache location
+# Token removed for security
+cache_dir = '/network/rit/lab/Lai_ReSecureAI/kiel/wmm'
+
+# Handle potential import conflicts with sentence_transformers
+try:
+    # Try to import bert_score directly to avoid sentence_transformers conflicts
+    from bert_score import score as bert_score
+    SIMILARITY_AVAILABLE = True
+    
+    def calc_scores_bert(original_sentence, substitute_sentences):
+        """BERTScore function using direct bert_score import."""
+        try:
+            # Safety check: truncate inputs if they're too long
+            max_chars = 2000  # Roughly 500 tokens
+            if len(original_sentence) > max_chars:
+                original_sentence = original_sentence[:max_chars]
+            
+            truncated_substitutes = []
+            for sub in substitute_sentences:
+                if len(sub) > max_chars:
+                    sub = sub[:max_chars]
+                truncated_substitutes.append(sub)
+            
+            references = [original_sentence] * len(truncated_substitutes)
+            P, R, F1 = bert_score(
+                cands=truncated_substitutes, 
+                refs=references, 
+                model_type="bert-base-uncased", 
+                verbose=False
+            )
+            return F1.tolist()
+        except Exception as e:
+            return [0.5] * len(substitute_sentences)
+    
+    def get_similarity_scores(original_sentence, substitute_sentences, method='bert'):
+        """Similarity function using direct bert_score import."""
+        if method == 'bert':
+            return calc_scores_bert(original_sentence, substitute_sentences)
+        else:
+            return [0.5] * len(substitute_sentences)
+            
+except ImportError as e:
+    print(f"Warning: bert_score import failed: {e}")
+    print("Falling back to neutral similarity scores...")
+    SIMILARITY_AVAILABLE = False
+    
+    def calc_scores_bert(original_sentence, substitute_sentences):
+        """Fallback BERTScore function with neutral scores."""
+        return [0.5] * len(substitute_sentences)
+    
+    def get_similarity_scores(original_sentence, substitute_sentences, method='bert'):
+        """Fallback similarity function with neutral scores."""
+        return [0.5] * len(substitute_sentences)
+
+# Setup NLTK data
+def setup_nltk_data():
+    """Setup NLTK data with error handling."""
+    try:
+        nltk.download('punkt_tab', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('averaged_perceptron_tagger_eng', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('wordnet', quiet=True)
+    except:
+        pass
+    try:
+        nltk.download('omw-1.4', quiet=True)
+    except:
+        pass
+
+setup_nltk_data()
+
+lemmatizer = WordNetLemmatizer()
+
+# Load spaCy model - download if not available
+try:
+    nlp = spacy.load("en_core_web_sm")
+except OSError:
+    print("Downloading spaCy model...")
+    subprocess.check_call([sys.executable, "-m", "spacy", "download", "en_core_web_sm"])
+    nlp = spacy.load("en_core_web_sm")
+
+# Define the detailed whitelist of POS tags (excluding adverbs)
+DETAILED_POS_WHITELIST = {
+    'NN',  # Noun, singular or mass (e.g., dog, car)
+    'NNS', # Noun, plural (e.g., dogs, cars)
+    'VB',  # Verb, base form (e.g., run, eat)
+    'VBD', # Verb, past tense (e.g., ran, ate)
+    'VBG', # Verb, gerund or present participle (e.g., running, eating)
+    'VBN', # Verb, past participle (e.g., run, eaten)
+    'VBP', # Verb, non-3rd person singular present (e.g., run, eat)
+    'VBZ', # Verb, 3rd person singular present (e.g., runs, eats)
+    'JJ',  # Adjective (e.g., big, blue)
+    'JJR', # Adjective, comparative (e.g., bigger, bluer)
+    'JJS', # Adjective, superlative (e.g., biggest, bluest)
+    'RB',  # Adverb (e.g., very, silently)
+    'RBR', # Adverb, comparative (e.g., better)
+    'RBS'  # Adverb, superlative (e.g., best)
+}
+
+# Global caches for better performance
+_pos_cache = {}
+_antonym_cache = {}
+_word_validity_cache = {}
+
+def extract_entities_and_pos(text):
+    """
+    Detect eligible tokens for replacement while skipping:
+    - Named entities (e.g., names, locations, organizations).
+    - Compound words (e.g., "Opteron-based").
+    - Phrasal verbs (e.g., "make up", "focus on").
+    - Punctuation and non-POS-whitelisted tokens.
+    """
+    doc = nlp(text)
+    sentence_target_pairs = []  # List to hold (sentence, target word, token index)
+
+    for sent in doc.sents:
+        for token in sent:
+            # Skip named entities using token.ent_type_ (more reliable than a text match)
+            if token.ent_type_:
+                continue
+
+            # Skip standalone punctuation
+            if token.is_punct:
+                continue
+
+            # Skip compound words (e.g., "Opteron-based")
+            if "-" in token.text or token.dep_ in {"compound", "amod"}:
+                continue
+
+            # Skip phrasal verbs (e.g., "make up", "focus on")
+            if token.pos_ == "VERB" and any(child.dep_ == "prt" for child in token.children):
+                continue
+
+            # Include regular tokens matching the POS whitelist
+            if token.tag_ in DETAILED_POS_WHITELIST:
+                sentence_target_pairs.append((sent.text, token.text, token.i))
+
+    return sentence_target_pairs
+
+def preprocess_text(text):
+    """
+    Preprocesses the text to handle abbreviations, titles, and edge cases
+    where a period or other punctuation does not signify a sentence end.
+    Ensures figures, acronyms, and short names are left untouched.
+    """
+    # Protect common abbreviations like "U.S." and "Corp."
+    text = re.sub(r'\b(U\.S|U\.K|Corp|Inc|Ltd)\.', r'\1<PERIOD>', text)
+    
+    # Protect floating-point numbers or ranges like "3.57" or "1.48–2.10"
+    text = re.sub(r'(\b\d+)\.(\d+)', r'\1<PERIOD>\2', text)
+    
+    # Avoid modifying standalone single-letter initials in names (e.g., "J. Smith")
+    text = re.sub(r'\b([A-Z])\.(?=\s[A-Z])', r'\1<PERIOD>', text)
+
+    # Protect acronym-like patterns with dots, such as "F.B.I."
+    text = re.sub(r'\b([A-Z]\.){2,}[A-Z]\.', lambda m: m.group(0).replace('.', '<PERIOD>'), text)
+
+    return text
+
+def split_sentences(text):
+    """
+    Splits text into sentences while preserving original newlines exactly.
+    - Protects abbreviations, acronyms, and floating-point numbers.
+    - Only adds newlines where necessary without duplicating them.
+    """
+    # Step 1: Protect abbreviations, floating numbers, acronyms
+    text = re.sub(r'\b(U\.S\.|U\.K\.|Inc\.|Ltd\.|Corp\.|e\.g\.|i\.e\.|etc\.)\b', r'\1<ABBR>', text)
+    text = re.sub(r'(\b\d+)\.(\d+)', r'\1<FLOAT>\2', text)
+    text = re.sub(r'\b([A-Z]\.){2,}[A-Z]\.', lambda m: m.group(0).replace('.', '<ABBR>'), text)
+
+    # Step 2: Identify sentence boundaries without duplicating newlines
+    sentences = []
+    for line in text.splitlines(keepends=True):  # Retain original newlines
+        # Split only if punctuation marks end a sentence
+        split_line = re.split(r'(?<=[.!?])\s+', line.strip())
+        sentences.extend([segment + "\n" if line.endswith("\n") else segment for segment in split_line])
+
+    # Step 3: Restore protected patterns
+    return [sent.replace('<ABBR>', '.').replace('<FLOAT>', '.') for sent in sentences]
+
+@lru_cache(maxsize=10000)
+def is_valid_word(word):
+    """Check if a word is valid using WordNet (cached)."""
+    return bool(wn.synsets(word))
+
+@lru_cache(maxsize=5000)
+def get_word_pos_tags(word):
+    """Get POS tags for a word using both NLTK and spaCy (cached)."""
+    nltk_pos = nltk.pos_tag([word])[0][1]
+    spacy_pos = nlp(word)[0].pos_
+    return nltk_pos, spacy_pos
+
+@lru_cache(maxsize=5000)
+def get_word_lemma(word):
+    """Get lemmatized form of a word (cached)."""
+    return lemmatizer.lemmatize(word)
+
+@lru_cache(maxsize=2000)
+def get_word_antonyms(word):
+    """Get antonyms for a word (cached). Includes all lemmas from all synsets."""
+    target_synsets = wn.synsets(word)
+    antonyms = set()
+    
+    # Get antonyms from all synsets and all lemmas
+    for syn in target_synsets:
+        for lem in syn.lemmas():
+            for ant in lem.antonyms():
+                # Add the antonym word (first part before the dot)
+                antonyms.add(ant.name().split('.')[0])
+                # Also add other lemmas of the antonym for completeness
+                for alt_lem in wn.synsets(ant.name().split('.')[0]):
+                    for alt_ant_lem in alt_lem.lemmas():
+                        antonyms.add(alt_ant_lem.name().split('.')[0])
+    
+    return antonyms
+
+def _are_semantically_compatible(target, candidate):
+    """
+    Check if target and candidate are semantically compatible for replacement.
+    Returns False if they are specific nouns in the same category (e.g., different crops, fruits, animals).
+    """
+    try:
+        # Direct check: if target and candidate are both specific terms for crops, animals, etc.
+        # check if they're NOT near-synonyms
+        
+        # Agricultural/crop terms that shouldn't be swapped
+        agricultural_terms = ['soybean', 'corn', 'maize', 'wheat', 'rice', 'barley', 'oats', 'sorghum', 
+                             'millet', 'grain', 'cereal', 'pulse', 'bean', 'legume']
+        
+        # If both are agricultural terms and different, block
+        if (target.lower() in agricultural_terms and candidate.lower() in agricultural_terms and 
+            target.lower() != candidate.lower()):
+            return False
+        
+        target_synsets = wn.synsets(target)
+        cand_synsets = wn.synsets(candidate)
+        
+        if not target_synsets or not cand_synsets:
+            return True  # If no synsets, allow through
+        
+        # Check if they're near-synonyms (very similar) - if so, allow
+        # We can use path similarity to check if they're similar enough
+        max_similarity = 0.0
+        for t_syn in target_synsets:
+            for c_syn in cand_synsets:
+                try:
+                    similarity = t_syn.path_similarity(c_syn) or 0.0
+                    max_similarity = max(max_similarity, similarity)
+                except:
+                    pass
+        
+        # If they have high path similarity (>0.5), they're similar enough to allow
+        if max_similarity > 0.5:
+            return True
+        
+        # Otherwise, check if they share common direct hypernyms
+        target_hypernyms = set()
+        for syn in target_synsets:
+            # Get immediate hypernyms (parent concepts)
+            for hypernym in syn.hypernyms():
+                target_hypernyms.add(hypernym)
+        
+        cand_hypernyms = set()
+        for syn in cand_synsets:
+            for hypernym in syn.hypernyms():
+                cand_hypernyms.add(hypernym)
+        
+        # If they share hypernyms, check if they're both specific instances (not general terms)
+        common_hypernyms = target_hypernyms & cand_hypernyms
+        
+        if common_hypernyms:
+            # Check if both words are specific instances of the same category
+            # If so, they shouldn't be replaced with each other
+            # We identify this by checking if their hypernym has many siblings
+            for hypernym in common_hypernyms:
+                siblings = hypernym.hyponyms()
+                # If there are many specific instances (e.g., many crops, many fruits)
+                # it's likely a category with specific instances that shouldn't be interchanged
+                if len(siblings) > 3:
+                    # Check if hypernym name suggests a specific category
+                    hypernym_name = hypernym.name().split('.')[0]
+                    category_keywords = [
+                        'crop', 'grain', 'fruit', 'animal', 'bird', 'fish', 'company', 
+                        'country', 'city', 'brand', 'product', 'food', 'vehicle'
+                    ]
+                    
+                    # If the hypernym contains category keywords, these are likely
+                    # specific instances that shouldn't be swapped
+                    if any(keyword in hypernym_name for keyword in category_keywords):
+                        return False
+        
+        return True
+        
+    except Exception as e:
+        # On any error, allow the candidate through (conservative approach)
+        return True
+
+def create_context_windows(full_text, target_sentence, target_word, tokenizer, max_tokens=400):
+    """
+    Create context windows around the target sentence for better MLM generation.
+    Intelligently handles tokenizer length limits by preserving the most relevant context.
+    
+    Args:
+        full_text: The complete document text
+        target_sentence: The sentence containing the target word
+        target_word: The word to be replaced
+        tokenizer: The tokenizer to check length limits
+        max_tokens: Maximum tokens to use for context (leave room for instruction + mask)
+    
+    Returns:
+        List of context windows with different levels of context
+    """
+    # Split full text into sentences
+    sentences = split_sentences(full_text)
+    
+    # Find the target sentence index
+    target_sentence_idx = None
+    for i, sent in enumerate(sentences):
+        if target_sentence.strip() in sent.strip():
+            target_sentence_idx = i
+            break
+    
+    if target_sentence_idx is None:
+        return [target_sentence]  # Fallback to original sentence
+    
+    # Create context windows with sentence-prioritized approach
+    context_windows = []
+    
+    # Window 1: Just the target sentence (always include)
+    context_windows.append(target_sentence)
+    
+    # Window 2: Target sentence + 1 sentence before and after (if fits)
+    start_idx = max(0, target_sentence_idx - 1)
+    end_idx = min(len(sentences), target_sentence_idx + 2)
+    context_window = " ".join(sentences[start_idx:end_idx])
+    
+    try:
+        encoded_len = len(tokenizer.encode(context_window))
+        if encoded_len <= max_tokens:
+            context_windows.append(context_window)
+    except Exception as e:
+        pass
+    
+    # Window 3: Target sentence + 2 sentences before and after (if fits)
+    start_idx = max(0, target_sentence_idx - 2)
+    end_idx = min(len(sentences), target_sentence_idx + 3)
+    context_window = " ".join(sentences[start_idx:end_idx])
+    
+    try:
+        encoded_len = len(tokenizer.encode(context_window))
+        if encoded_len <= max_tokens:
+            context_windows.append(context_window)
+    except Exception as e:
+        pass
+    
+    # Window 4: Target sentence + 3 sentences before and after (if fits)
+    start_idx = max(0, target_sentence_idx - 3)
+    end_idx = min(len(sentences), target_sentence_idx + 4)
+    context_window = " ".join(sentences[start_idx:end_idx])
+    
+    try:
+        encoded_len = len(tokenizer.encode(context_window))
+        if encoded_len <= max_tokens:
+            context_windows.append(context_window)
+    except Exception as e:
+        pass
+    
+    # Window 5: Intelligent context with sentence prioritization + word expansion
+    intelligent_context = _create_intelligent_context(
+        full_text, target_word, target_sentence_idx, tokenizer, max_tokens
+    )
+    context_windows.append(intelligent_context)
+    
+    return context_windows
+
+def _create_intelligent_context(full_text, target_word, target_sentence_idx, tokenizer, max_tokens):
+    """
+    Create intelligent context that prioritizes sentence boundaries while respecting token limits.
+    Strategy: Target sentence → Nearby sentences → Word-level expansion
+    """
+    sentences = split_sentences(full_text)
+    
+    # Strategy 1: Always start with the target sentence
+    target_sentence = sentences[target_sentence_idx]
+    try:
+        target_sentence_tokens = len(tokenizer.encode(target_sentence))
+    except Exception as e:
+        target_sentence_tokens = 1000  # Fallback to assume it's too long
+    
+    if target_sentence_tokens > max_tokens:
+        # If even target sentence is too long, truncate intelligently
+        return _truncate_sentence_intelligently(target_sentence, target_word, tokenizer, max_tokens)
+    
+    # Strategy 2: Expand sentence-by-sentence around target sentence
+    best_context = target_sentence
+    best_token_count = target_sentence_tokens
+    
+    # Try adding sentences before and after the target sentence
+    for sentence_radius in range(1, min(len(sentences), 20)):  # Max 20 sentences radius
+        start_idx = max(0, target_sentence_idx - sentence_radius)
+        end_idx = min(len(sentences), target_sentence_idx + sentence_radius + 1)
+        
+        # Create context with complete sentences
+        context_sentences = sentences[start_idx:end_idx]
+        context_window = " ".join(context_sentences)
+        try:
+            token_count = len(tokenizer.encode(context_window))
+        except Exception as e:
+            token_count = 1000  # Fallback to assume it's too long
+        
+        if token_count <= max_tokens:
+            # This sentence expansion fits, keep it as our best option
+            best_context = context_window
+            best_token_count = token_count
+        else:
+            # This expansion is too big, stop here
+            break
+    
+    # Strategy 3: If we have room left, try word-level expansion within the best sentence context
+    remaining_tokens = max_tokens - best_token_count
+    if remaining_tokens > 50:  # If we have significant room left
+        enhanced_context = _enhance_with_word_expansion(
+            full_text, target_word, best_context, tokenizer, remaining_tokens
+        )
+        if enhanced_context:
+            return enhanced_context
+    
+    return best_context
+
+def _enhance_with_word_expansion(full_text, target_word, current_context, tokenizer, remaining_tokens):
+    """
+    Enhance the current sentence-based context with word-level expansion if there's room.
+    """
+    words = full_text.split()
+    target_word_idx = None
+    
+    # Find target word position in full text
+    for i, word in enumerate(words):
+        if word.lower() == target_word.lower():
+            target_word_idx = i
+            break
+    
+    if target_word_idx is None:
+        return current_context
+    
+    # Try to expand word-by-word around the target word
+    try:
+        current_tokens = len(tokenizer.encode(current_context))
+    except Exception as e:
+        print(f"WARNING: Error encoding current context: {e}")
+        current_tokens = 1000  # Fallback to assume it's too long
+    
+    for expansion_size in range(1, min(len(words), 100)):  # Max 100 words expansion
+        start_word = max(0, target_word_idx - expansion_size)
+        end_word = min(len(words), target_word_idx + expansion_size + 1)
+        
+        expanded_context = " ".join(words[start_word:end_word])
+        try:
+            expanded_tokens = len(tokenizer.encode(expanded_context))
+        except Exception as e:
+            expanded_tokens = 1000  # Fallback to assume it's too long
+        
+        if expanded_tokens <= current_tokens + remaining_tokens:
+            # This expansion fits within our remaining token budget
+            return expanded_context
+        else:
+            # This expansion is too big, stop here
+            break
+    
+    return current_context
+
+def _truncate_sentence_intelligently(sentence, target_word, tokenizer, max_tokens):
+    """
+    Intelligently truncate a sentence while preserving context around the target word.
+    """
+    words = sentence.split()
+    target_word_idx = None
+    
+    # Find target word position
+    for i, word in enumerate(words):
+        if word.lower() == target_word.lower():
+            target_word_idx = i
+            break
+    
+    if target_word_idx is None:
+        # If target word not found, truncate from the end
+        truncated = " ".join(words)
+        try:
+            while len(tokenizer.encode(truncated)) > max_tokens and len(words) > 1:
+                words = words[:-1]
+                truncated = " ".join(words)
+        except Exception as e:
+            # Fallback: return first few words
+            truncated = " ".join(words[:10]) if len(words) >= 10 else " ".join(words)
+        return truncated
+    
+    # Truncate symmetrically around target word
+    context_words = 10  # Start with 10 words before/after
+    while context_words > 0:
+        start_word = max(0, target_word_idx - context_words)
+        end_word = min(len(words), target_word_idx + context_words + 1)
+        truncated_sentence = " ".join(words[start_word:end_word])
+        
+        try:
+            if len(tokenizer.encode(truncated_sentence)) <= max_tokens:
+                return truncated_sentence
+        except Exception as e:
+            # Continue to next iteration
+            pass
+        
+        context_words -= 1
+    
+    # Fallback: just the target word with minimal context
+    return f"... {target_word} ..."
+
+def _intelligent_token_slicing(input_text, tokenizer, max_length=512, mask_token_id=None):
+    """
+    Intelligently slice input text to fit within max_length tokens while preserving the mask token.
+    Strategy: Preserve mask token and surrounding context, remove excess tokens from less important areas.
+    
+    Args:
+        input_text: The full input text to be tokenized
+        tokenizer: The tokenizer to use
+        max_length: Maximum allowed sequence length (default 512)
+        mask_token_id: The mask token ID to preserve
+    
+    Returns:
+        Tuple of (sliced_input_ids, mask_position_in_sliced)
+    """
+    # First, tokenize the full input
+    input_ids = tokenizer.encode(input_text, add_special_tokens=True)
+    
+    # If already within limits, return as is
+    if len(input_ids) <= max_length:
+        mask_pos = input_ids.index(mask_token_id) if mask_token_id in input_ids else None
+        return input_ids, mask_pos
+    
+    # Find mask token position
+    mask_positions = [i for i, token_id in enumerate(input_ids) if token_id == mask_token_id]
+    
+    if not mask_positions:
+        # No mask token found, truncate from the end
+        return input_ids[:max_length], None
+    
+    mask_pos = mask_positions[0]  # Use first mask token
+    
+    # Calculate how many tokens we need to remove
+    excess_tokens = len(input_ids) - max_length
+    
+    # Strategy: Remove tokens from both ends while preserving mask context
+    # Reserve some context around the mask token
+    mask_context_size = min(50, max_length // 4)  # Reserve 25% of max_length or 50 tokens, whichever is smaller
+    
+    # Calculate available space for context around mask
+    available_before = min(mask_pos, mask_context_size)
+    available_after = min(len(input_ids) - mask_pos - 1, mask_context_size)
+    
+    # Calculate how much to remove from each end
+    tokens_to_remove_before = max(0, mask_pos - available_before)
+    tokens_to_remove_after = max(0, (len(input_ids) - mask_pos - 1) - available_after)
+    
+    # Initialize removal variables
+    remove_before = 0
+    remove_after = 0
+    
+    # Distribute excess tokens proportionally
+    if excess_tokens > 0:
+        if tokens_to_remove_before + tokens_to_remove_after >= excess_tokens:
+            # We can remove enough from the ends
+            if tokens_to_remove_before >= excess_tokens // 2:
+                remove_before = excess_tokens // 2
+                remove_after = excess_tokens - remove_before
+            else:
+                remove_before = tokens_to_remove_before
+                remove_after = min(tokens_to_remove_after, excess_tokens - remove_before)
+        else:
+            # Need to remove more aggressively
+            remove_before = tokens_to_remove_before
+            remove_after = tokens_to_remove_after
+            remaining_excess = excess_tokens - remove_before - remove_after
+            
+            # Remove remaining excess from the end
+            if remaining_excess > 0:
+                remove_after += remaining_excess
+    
+    # Calculate final indices
+    start_idx = remove_before
+    end_idx = len(input_ids) - remove_after
+    
+    # Ensure we don't exceed max_length
+    if end_idx - start_idx > max_length:
+        # Center around mask token
+        half_length = max_length // 2
+        start_idx = max(0, mask_pos - half_length)
+        end_idx = min(len(input_ids), start_idx + max_length)
+    
+    # Slice the input_ids
+    sliced_input_ids = input_ids[start_idx:end_idx]
+    
+    # Debug information
+    if len(sliced_input_ids) > max_length:
+        # Force truncation as final fallback
+        sliced_input_ids = sliced_input_ids[:max_length]
+    
+    # Adjust mask position for the sliced sequence
+    adjusted_mask_pos = mask_pos - start_idx
+    
+    return sliced_input_ids, adjusted_mask_pos
+
+def _create_word_level_context(full_text, target_word, tokenizer, max_tokens):
+    """
+    Create context by expanding word-by-word around the target word until reaching token limit.
+    This maximizes context while respecting tokenizer limits.
+    """
+    words = full_text.split()
+    target_word_idx = None
+    
+    # Find target word position in full text
+    for i, word in enumerate(words):
+        if word.lower() == target_word.lower():
+            target_word_idx = i
+            break
+    
+    if target_word_idx is None:
+        # Fallback: expand from beginning until token limit
+        return _expand_from_start(words, tokenizer, max_tokens)
+    
+    # Word-by-word expansion around target word
+    return _expand_around_target(words, target_word_idx, tokenizer, max_tokens)
+
+def _expand_around_target(words, target_idx, tokenizer, max_tokens):
+    """
+    Expand word-by-word around target word until reaching token limit.
+    """
+    best_context = ""
+    best_token_count = 0
+    
+    # Try different expansion sizes
+    for expansion_size in range(1, min(len(words), 200)):  # Max 200 words expansion
+        start_word = max(0, target_idx - expansion_size)
+        end_word = min(len(words), target_idx + expansion_size + 1)
+        
+        context_window = " ".join(words[start_word:end_word])
+        try:
+            token_count = len(tokenizer.encode(context_window))
+        except Exception as e:
+            token_count = 1000  # Fallback to assume it's too long
+        
+        if token_count <= max_tokens:
+            # This expansion fits, keep it as our best option
+            best_context = context_window
+            best_token_count = token_count
+        else:
+            # This expansion is too big, stop here
+            break
+    
+    # If we found a good context, return it
+    if best_context:
+        return best_context
+    
+    # Fallback: minimal context around target word
+    start_word = max(0, target_idx - 5)
+    end_word = min(len(words), target_idx + 6)
+    return " ".join(words[start_word:end_word])
+
+def _expand_from_start(words, tokenizer, max_tokens):
+    """
+    Expand from the start of the text until reaching token limit.
+    """
+    for end_idx in range(len(words), 0, -1):
+        context_window = " ".join(words[:end_idx])
+        try:
+            if len(tokenizer.encode(context_window)) <= max_tokens:
+                return context_window
+        except Exception as e:
+            # Continue to next iteration
+            pass
+    
+    # Fallback: first few words
+    return " ".join(words[:10]) if len(words) >= 10 else " ".join(words)
+
+def whole_context_mlm_inference(full_text, sentence_target_pairs, tokenizer, lm_model, Top_K=20, batch_size=32, max_context_tokens=400, max_length=512, similarity_context_mode='whole'):
+    """
+    Enhanced MLM inference using whole document context for better candidate generation.
+    """
+    results = {}
+    
+    # Group targets by sentence for batch processing
+    sentence_groups = {}
+    for sent, target, index in sentence_target_pairs:
+        if sent not in sentence_groups:
+            sentence_groups[sent] = []
+        sentence_groups[sent].append((target, index))
+    
+    for sentence, targets in sentence_groups.items():
+        # Process targets in batches
+        for i in range(0, len(targets), batch_size):
+            batch_targets = targets[i:i+batch_size]
+            batch_results = _process_whole_context_mlm_batch(
+                full_text, sentence, batch_targets, tokenizer, lm_model, Top_K, max_context_tokens, max_length, similarity_context_mode
+            )
+            results.update(batch_results)
+    
+    return results
+
+def _process_whole_context_mlm_batch(full_text, sentence, targets, tokenizer, lm_model, Top_K, max_context_tokens=400, max_length=512, similarity_context_mode='whole'):
+    """
+    Process a batch of targets using whole document context for MLM.
+    """
+    results = {}
+    
+    # Tokenize sentence once
+    doc = nlp(sentence)
+    tokens = [token.text for token in doc]
+    
+    # Create multiple masked versions for batch processing
+    masked_inputs = []
+    mask_positions = []
+    contexts_for_targets = []
+    
+    for target, index in targets:
+        if index < len(tokens):
+            # Create context windows with tokenizer length awareness
+            context_windows = create_context_windows(full_text, sentence, target, tokenizer, max_tokens=max_context_tokens)
+            
+            # Use the most comprehensive context window that fits within token limits
+            full_context = context_windows[-1]  # Built around the target sentence
+            # Select context for similarity according to mode
+            context = sentence if similarity_context_mode == 'sentence' else full_context
+            
+            # Create masked version of the FULL context (not just the sentence)
+            masked_full_context = context.replace(target, tokenizer.mask_token, 1)
+            
+            instruction = "Given the full document context, replace the masked word with a word that fits grammatically, preserves the original meaning, and ensures natural flow in the document:"
+            input_text = f"{instruction} {context} {tokenizer.sep_token} {masked_full_context}"
+            
+            # AGGRESSIVE FIX: Truncate input text BEFORE tokenization to prevent errors
+            # Estimate token count (roughly 1 token per 4 characters for English)
+            estimated_tokens = len(input_text) // 4
+            if estimated_tokens > 500:  # Leave some buffer
+                # Truncate to roughly 2000 characters (500 tokens)
+                input_text = input_text[:2000]
+            
+            # SIMPLE FIX: Truncate input text if it's too long
+            try:
+                temp_tokens = tokenizer.encode(input_text, add_special_tokens=True)
+                if len(temp_tokens) > 512:
+                    # Truncate the input text by removing words from the end
+                    words = input_text.split()
+                    while len(tokenizer.encode(" ".join(words), add_special_tokens=True)) > 512 and len(words) > 10:
+                        words = words[:-1]
+                    input_text = " ".join(words)
+            except Exception as e:
+                # Emergency truncation - just take first 200 words
+                words = input_text.split()
+                input_text = " ".join(words[:200])
+            
+            masked_inputs.append(input_text)
+            # Store the original sentence-level index for reference, but mask position will be calculated during tokenization
+            mask_positions.append(index)
+            contexts_for_targets.append(context)
+    
+    if not masked_inputs:
+        return results
+    
+    # Batch tokenize
+    MAX_LENGTH = max_length  # Use parameter for A100 optimization
+    batch_inputs = []
+    batch_mask_positions = []
+    batch_contexts = []
+    
+    for input_text, mask_pos in zip(masked_inputs, mask_positions):
+        # Use intelligent token slicing to ensure we stay within MAX_LENGTH
+        try:
+            input_ids, adjusted_mask_pos = _intelligent_token_slicing(
+                input_text, tokenizer, max_length=MAX_LENGTH, mask_token_id=tokenizer.mask_token_id
+            )
+            
+            if adjusted_mask_pos is not None:
+                batch_inputs.append(input_ids)
+                batch_mask_positions.append(adjusted_mask_pos)
+            else:
+                # Mask token not found in sliced sequence, skip this input
+                continue
+                
+        except Exception as e:
+            # Fallback: simple truncation
+            try:
+                input_ids = tokenizer.encode(input_text, add_special_tokens=True)
+                if len(input_ids) > MAX_LENGTH:
+                    input_ids = input_ids[:MAX_LENGTH]
+                
+                masked_position = input_ids.index(tokenizer.mask_token_id)
+                batch_inputs.append(input_ids)
+                batch_mask_positions.append(masked_position)
+            except ValueError:
+                # Mask token not found, skip this input
+                continue
+    
+    if not batch_inputs:
+        return results
+    
+    # Pad sequences to same length, but ensure we don't exceed MAX_LENGTH
+    max_len = min(max(len(ids) for ids in batch_inputs), MAX_LENGTH)
+    
+    # Additional safety check: truncate any sequences that are still too long
+    truncated_batch_inputs = []
+    for input_ids in batch_inputs:
+        if len(input_ids) > MAX_LENGTH:
+            input_ids = input_ids[:MAX_LENGTH]
+        truncated_batch_inputs.append(input_ids)
+    
+    padded_inputs = []
+    attention_masks = []
+    
+    for input_ids in truncated_batch_inputs:
+        attention_mask = [1] * len(input_ids) + [0] * (max_len - len(input_ids))
+        padded_ids = input_ids + [tokenizer.pad_token_id] * (max_len - len(input_ids))
+        padded_inputs.append(padded_ids)
+        attention_masks.append(attention_mask)
+    
+    # Final safety check: ensure no sequence exceeds MAX_LENGTH
+    for i, padded_ids in enumerate(padded_inputs):
+        if len(padded_ids) > MAX_LENGTH:
+            padded_inputs[i] = padded_ids[:MAX_LENGTH]
+            attention_masks[i] = attention_masks[i][:MAX_LENGTH]
+    
+    # Batch inference - optimized for A100 with mixed precision
+    with torch.no_grad():
+        input_tensor = torch.tensor(padded_inputs, dtype=torch.long)
+        attention_tensor = torch.tensor(attention_masks, dtype=torch.long)
+        
+        if torch.cuda.is_available():
+            input_tensor = input_tensor.cuda()
+            attention_tensor = attention_tensor.cuda()
+            
+            # Use mixed precision for A100 optimization
+            with torch.amp.autocast('cuda'):
+                outputs = lm_model(input_tensor, attention_mask=attention_tensor)
+                logits = outputs.logits
+        else:
+            outputs = lm_model(input_tensor, attention_mask=attention_tensor)
+            logits = outputs.logits
+    
+    # Process results - collect filtered candidates first
+    batch_filtered_results = {}
+    for i, (target, index) in enumerate(targets):
+        if i < len(batch_mask_positions):
+            mask_pos = batch_mask_positions[i]
+            mask_logits = logits[i, mask_pos].squeeze()
+            
+            # Get top predictions
+            top_tokens = torch.topk(mask_logits, k=Top_K, dim=0)[1]
+            scores = torch.softmax(mask_logits, dim=0)[top_tokens].tolist()
+            words = [tokenizer.decode(token.item()).strip() for token in top_tokens]
+            
+            # Filter candidates (without similarity scoring)
+            filtered_candidates = _filter_candidates_batch(target, words, scores, tokens, index)
+            if filtered_candidates:
+                # Attach the exact context window used for this target
+                batch_filtered_results[(sentence, target, index)] = {
+                    'filtered_words': filtered_candidates,
+                    'context': contexts_for_targets[i]
+                }
+    
+    # Batch similarity scoring for all candidates
+    if batch_filtered_results:
+        similarity_results = _batch_similarity_scoring(batch_filtered_results, tokenizer)
+        results.update(similarity_results)
+    
+    return results
+
+def _filter_candidates_batch(target, words, scores, tokens, target_index):
+    """
+    Optimized batch filtering of candidates (no similarity scoring here - moved to batch level).
+    """
+    # Basic filtering
+    filtered_words = []
+    filtered_scores = []
+    seen_words = set()
+    
+    for word, score in zip(words, scores):
+        word_lower = word.lower()
+        if word_lower in seen_words or word_lower == target.lower():
+            continue
+        seen_words.add(word_lower)
+        
+        if not is_valid_word(word):
+            continue
+        
+        # Quick POS check
+        target_nltk_pos, target_spacy_pos = get_word_pos_tags(target)
+        cand_nltk_pos, cand_spacy_pos = get_word_pos_tags(word)
+        
+        if target_nltk_pos != cand_nltk_pos or target_spacy_pos != cand_spacy_pos:
+            continue
+        
+        # Check antonyms (bidirectional and case-insensitive)
+        antonyms = get_word_antonyms(target)
+        if word.lower() in [ant.lower() for ant in antonyms]:
+            continue
+        
+        # Also check if the candidate has the target as an antonym (reverse check)
+        candidate_antonyms = get_word_antonyms(word)
+        if target.lower() in [ant.lower() for ant in candidate_antonyms]:
+            continue
+        
+        # Hardcoded common antonym pairs (for words not in WordNet or as additional safeguard)
+        common_antonyms = {
+            'big': ['small', 'tiny', 'little'],
+            'small': ['big', 'large', 'huge'],
+            'large': ['small', 'tiny', 'little'],
+            'good': ['bad', 'evil', 'wrong'],
+            'bad': ['good', 'great', 'excellent'],
+            'high': ['low'],
+            'low': ['high'],
+            'new': ['old'],
+            'old': ['new'],
+            'fast': ['slow'],
+            'slow': ['fast'],
+            'rich': ['poor'],
+            'poor': ['rich'],
+            'hot': ['cold'],
+            'cold': ['hot'],
+            'happy': ['sad', 'unhappy'],
+            'sad': ['happy', 'joyful'],
+            'true': ['false', 'untrue'],
+            'false': ['true'],
+            'real': ['fake', 'unreal'],
+            'fake': ['real'],
+            'up': ['down'],
+            'down': ['up'],
+            'yes': ['no'],
+            'no': ['yes'],
+            'alive': ['dead'],
+            'dead': ['alive'],
+            'safe': ['unsafe', 'dangerous'],
+            'dangerous': ['safe'],
+            'clean': ['dirty'],
+            'dirty': ['clean'],
+            'full': ['empty'],
+            'empty': ['full'],
+            'open': ['closed', 'shut'],
+            'closed': ['open'],
+            'begin': ['end', 'finish'],
+            'end': ['begin', 'start'],
+            'start': ['end', 'finish'],
+            'finish': ['start', 'begin'],
+            'first': ['last'],
+            'last': ['first']
+        }
+        
+        # Check if word is a known antonym of target (case-insensitive)
+        target_lower = target.lower()
+        if target_lower in common_antonyms and word.lower() in common_antonyms[target_lower]:
+            continue
+        
+        # Check if word and target are in the same specific noun category (e.g., crops, animals, companies)
+        # If they are different specific terms in the same category, exclude the candidate
+        if not _are_semantically_compatible(target, word):
+            continue
+        
+        filtered_words.append(word)
+        filtered_scores.append(score)
+    
+    if len(filtered_words) < 2:
+        return None
+    
+    # Return filtered words without similarity scoring (done at batch level)
+    return filtered_words
+
+def _batch_similarity_scoring(batch_results, tokenizer):
+    """
+    Optimized batched similarity scoring across multiple sentences for full context.
+    Processes all candidates from multiple sentences together for better efficiency.
+    """
+    # Collect all similarity scoring tasks
+    similarity_tasks = []
+    sentence_contexts = {}
+    
+    for (sentence, target, index), value in batch_results.items():
+        if value is None:
+            continue
+        # Support both legacy list and new dict with context
+        if isinstance(value, dict):
+            filtered_words = value.get('filtered_words')
+            context = value.get('context', sentence)
+        else:
+            filtered_words = value
+            context = sentence
+            
+        # Tokenize the sentence once
+        tokens = tokenizer.tokenize(sentence)
+        if index >= len(tokens):
+            continue
+            
+        # Store sentence context for later use
+        sentence_contexts[(sentence, target, index)] = {
+            'tokens': tokens,
+            'target_index': index,
+            'filtered_words': filtered_words
+        }
+        
+        # Create candidate sentences for this target
+        for word in filtered_words:
+            candidate_tokens = tokens.copy()
+            candidate_tokens[index] = word
+            candidate_sentence = tokenizer.convert_tokens_to_string(candidate_tokens)
+            
+            # Build full-context candidate by replacing the sentence inside the chosen context once
+            candidate_full_context = context.replace(sentence, candidate_sentence, 1)
+            similarity_tasks.append({
+                'original_context': context,
+                'candidate_full_context': candidate_full_context,
+                'target_word': word,
+                'context_key': (sentence, target, index)
+            })
+    
+    if not similarity_tasks:
+        return {}
+    
+    # Batch process all similarity scoring
+    try:
+        # Group by original full context for efficient BERTScore computation
+        context_groups = {}
+        for task in similarity_tasks:
+            orig_ctx = task['original_context']
+            if orig_ctx not in context_groups:
+                context_groups[orig_ctx] = []
+            context_groups[orig_ctx].append(task)
+        
+        # Process each context group
+        final_results = {}
+        for orig_context, tasks in context_groups.items():
+            # Extract candidate full-contexts
+            candidate_contexts = [task['candidate_full_context'] for task in tasks]
+            
+            # Batch BERTScore computation against the same full context
+            try:
+                similarity_scores = calc_scores_bert(orig_context, candidate_contexts)
+            except Exception as e:
+                # Fallback to neutral scores
+                similarity_scores = [0.5] * len(candidate_contexts)
+            
+            if similarity_scores and not all(score == 0.5 for score in similarity_scores):
+                # Group results by context key
+                for task, score in zip(tasks, similarity_scores):
+                    context_key = task['context_key']
+                    if context_key not in final_results:
+                        final_results[context_key] = []
+                    final_results[context_key].append((task['target_word'], score))
+        
+        # Sort results by similarity score
+        for context_key in final_results:
+            final_results[context_key].sort(key=lambda x: x[1], reverse=True)
+        
+        return final_results
+        
+    except Exception as e:
+        return {}
+
+def parallel_tournament_sampling(target_results, secret_key, m, c, h, alpha):
+    """
+    Parallel tournament sampling for multiple targets.
+    """
+    results = {}
+    
+    if not target_results:
+        return results
+    
+    def process_single_tournament(item):
+        (sentence, target, index), candidates = item
+        if not candidates:
+            return (sentence, target, index), None
+        
+        alternatives = [alt[0] for alt in candidates]
+        similarity = [alt[1] for alt in candidates]
+        
+        if not alternatives or not similarity:
+            return (sentence, target, index), None
+        
+        # Get context
+        context_tokens = word_tokenize(sentence)
+        left_context = context_tokens[max(0, index - h):index]
+        
+        # Tournament selection
+        from SynthID_randomization import tournament_select_word
+        randomized_word = tournament_select_word(
+            target, alternatives, similarity, 
+            context=left_context, key=secret_key, m=m, c=c, alpha=alpha
+        )
+        
+        return (sentence, target, index), randomized_word
+    
+    # Process in parallel
+    max_workers = max(1, min(8, len(target_results)))
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_item = {executor.submit(process_single_tournament, item): item for item in target_results.items()}
+        
+        for future in as_completed(future_to_item):
+            key, result = future.result()
+            results[key] = result
+    
+    return results
+
+def whole_context_process_sentence(full_text, sentence, tokenizer, lm_model, Top_K, threshold, secret_key, m, c, h, alpha, output_name, batch_size=32, max_length=512, max_context_tokens=400, similarity_context_mode='whole'):
+    """
+    Enhanced sentence processing using whole document context for better candidate generation.
+    """
+    replacements = []
+    sampling_results = []
+    doc = nlp(sentence)
+    sentence_target_pairs = extract_entities_and_pos(sentence)
+
+    if not sentence_target_pairs:
+        return replacements, sampling_results
+
+    # Filter valid target pairs
+    valid_pairs = []
+    spacy_tokens = [token.text for token in doc]
+    
+    for sent, target, position in sentence_target_pairs:
+        if position < len(spacy_tokens) and spacy_tokens[position] == target:
+            valid_pairs.append((sent, target, position))
+
+    if not valid_pairs:
+        return replacements, sampling_results
+
+    # Enhanced MLM inference with whole document context
+    batch_results = whole_context_mlm_inference(full_text, valid_pairs, tokenizer, lm_model, Top_K, batch_size, max_context_tokens, max_length, similarity_context_mode)
+    
+    # Filter by threshold (matching original logic)
+    filtered_results = {}
+    for key, candidates in batch_results.items():
+        if candidates:
+            # Apply threshold filtering (matching original logic)
+            threshold_candidates = [(word, score) for word, score in candidates if score >= threshold]
+            if len(threshold_candidates) >= 2:
+                filtered_results[key] = threshold_candidates
+    
+    # Parallel tournament sampling
+    tournament_results = parallel_tournament_sampling(filtered_results, secret_key, m, c, h, alpha)
+    
+    # Collect replacements and sampling results
+    for (sent, target, position), randomized_word in tournament_results.items():
+        if randomized_word:
+            # Get the alternatives for this target from the filtered results
+            alternatives = filtered_results.get((sent, target, position), [])
+            alternatives_list = [alt[0] for alt in alternatives]
+            # Include similarity scores for each alternative (preserve old 'alternatives' list for compatibility)
+            alternatives_with_similarity = [
+                {"word": alt[0], "similarity": float(alt[1])} for alt in alternatives
+            ]
+            
+            # Track sampling results
+            sampling_results.append({
+                "word": target,
+                "alternatives": alternatives_list,
+                "alternatives_with_similarity": alternatives_with_similarity,
+                "randomized_word": randomized_word
+            })
+            
+            replacements.append((position, target, randomized_word))
+    
+    return replacements, sampling_results
+
+# Legacy function for compatibility
+def look_up(sentence, target, index, tokenizer, lm_model, Top_K=20, threshold=0.75):
+    """
+    Legacy single-target lookup function for compatibility.
+    """
+    results = batch_mlm_inference([(sentence, target, index)], tokenizer, lm_model, Top_K)
+    return results.get((sentence, target, index), None)
+
+def batch_mlm_inference(sentence_target_pairs, tokenizer, lm_model, Top_K=20):
+    """
+    Legacy batch MLM inference function for compatibility.
+    """
+    return whole_context_mlm_inference("", sentence_target_pairs, tokenizer, lm_model, Top_K)
+
+def batch_look_up(sentence_target_pairs, tokenizer, lm_model, Top_K=20, threshold=0.75, max_workers=4):
+    """
+    Optimized batch lookup using the new batch MLM inference.
+    """
+    return batch_mlm_inference(sentence_target_pairs, tokenizer, lm_model, Top_K)