inintal commit

.dockerignore

@@ -0,0 +1,20 @@
+__pycache__
+*.pyc
+*.pyo
+*.pyd
+.Python
+env
+pip-log.txt
+pip-delete-this-directory.txt
+.tox
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.log
+.pytest_cache
+.env
+.venv
+.DS_Store

.gitattributes

@@ -1,6 +1,8 @@
+*.pdf filter=lfs diff=lfs merge=lfs -text
+*.txt filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
 *.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
@@ -33,3 +35,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+static/ia_gen_droits_auteur.pdf filter=lfs diff=lfs merge=lfs -text
+wm_interactive/static/hf_cache/** filter=lfs diff=lfs merge=lfs -text
+wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/5af571cbf074e6d21a03528d2330792e532ca608f24ac70a143f6b369968ab8c filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+.DS_Store

Dockerfile

@@ -0,0 +1,24 @@
+FROM python:3.9-slim
+
+WORKDIR /app
+
+# Copy only the requirements first to leverage Docker cache
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the rest of the application
+COPY wm_interactive/ ./wm_interactive/
+COPY run.py .
+
+# Create necessary directories
+RUN mkdir -p wm_interactive/static/hf_cache
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV FLASK_APP=run.py
+
+# Expose the port the app runs on
+EXPOSE 7860
+
+# Command to run the application
+CMD ["python", "run.py"]

README.md

@@ -1,12 +1,14 @@
 ---
-title: Interactive Llm Wm
-emoji: 😻
-colorFrom: green
-colorTo: blue
+title: Interactive Text Watermark Detection
+emoji: 📝
+colorFrom: blue
+colorTo: pink
 sdk: docker
 pinned: false
 license: apache-2.0
-short_description: An interactive demo for LLM watermarking
+short_description: An interactive demo for detection of text watermarks
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Interactive Text Watermark Detection
+
+This repository contains the code for an interactive demo for detection of watermarked text generated from LLM (Large Language Model) models.

data/prompts.json

@@ -0,0 +1,52 @@
+[
+    {
+        "instruction": "Write a short story about a robot learning to paint.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Explain how photosynthesis works in simple terms.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Write a recipe for chocolate chip cookies.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Describe the main differences between classical and quantum computing.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Write a haiku about the changing seasons.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Explain why the sky appears blue during the day.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Write a short dialogue between two friends discussing their favorite books.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Describe three ways to reduce your carbon footprint.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Write a brief explanation of how the internet works.",
+        "input": "",
+        "output": ""
+    },
+    {
+        "instruction": "Create a short motivational speech about perseverance.",
+        "input": "",
+        "output": ""
+    }
+]

requirements.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:59ef79cf6a8a998de982ccc64e93ca2b7602aa989b38b2c264d385acf728ef80
+size 75

run.py

@@ -0,0 +1,12 @@
+"""
+Main entry point for the watermark detection application.
+Run with: python run.py
+
+docker build -t wm-interactive .
+docker run -p 7860:7860 wm-interactive
+"""
+
+from wm_interactive.web.app import app
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860)

sandbox.ipynb

@@ -0,0 +1,81 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from transformers import AutoTokenizer, LlamaForCausalLM\n",
+    "\n",
+    "model_id = \"meta-llama/Llama-3.2-1B-Instruct\"\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=\"wm_detector/static/hf_cache\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "[4438, 311, 1304, 264, 19692]\n",
+      "['How', 'Ġto', 'Ġmake', 'Ġa', 'Ġcake']\n",
+      "['How', ' to', ' make', ' a', ' cake']\n"
+     ]
+    }
+   ],
+   "source": [
+    "def tokenize_text(text):\n",
+    "    return tokenizer.encode(text, add_special_tokens=False)\n",
+    "\n",
+    "text = \"How to make a cake\"\n",
+    "token_ids = tokenize_text(text)\n",
+    "tokens = tokenizer.convert_ids_to_tokens(token_ids)\n",
+    "token_strs = [tokenizer.convert_tokens_to_string([token]) for token in tokens]\n",
+    "decoded = tokenizer.decode(tokenize_text(text))\n",
+    "\n",
+    "print(token_ids)\n",
+    "print(tokens)\n",
+    "print(token_strs)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "base",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.2"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

wm_interactive/core/detector.py

@@ -0,0 +1,263 @@
+
+import numpy as np
+from scipy import special
+
+import torch
+from transformers import AutoTokenizer
+
+from .hashing import get_seed_rng
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class WmDetector():
+    def __init__(self, 
+        tokenizer: AutoTokenizer, 
+        ngram: int = 1,
+        seed: int = 0
+    ):
+        # model config
+        self.tokenizer = tokenizer
+        self.vocab_size = self.tokenizer.vocab_size
+        # watermark config
+        self.ngram = ngram
+        self.seed = seed
+        self.rng = torch.Generator()
+        self.rng.manual_seed(self.seed)
+
+    def aggregate_scores(
+        self, 
+        scores: list[np.array], 
+        aggregation: str = 'mean'
+    ) -> float:
+        """Aggregate scores along a text."""
+        if aggregation == 'sum':
+           return scores.sum(axis=0)
+        elif aggregation == 'mean':
+            return scores.mean(axis=0)
+        elif aggregation == 'max':
+            return scores.max(axis=0)
+        else:
+             raise ValueError(f'Aggregation {aggregation} not supported.')
+
+    def get_details(
+        self, 
+        text: str,
+        scoring_method: str="v2",
+        ntoks_max: int = None,
+    ) -> list[dict]:
+        """
+        Get score increment for each token in text.
+        Args:
+            text: input text
+            scoring_method: 
+                'none': score all ngrams
+                'v1': only score tokens for which wm window is unique
+                'v2': only score unique {wm window+tok} is unique
+            ntoks_max: maximum number of tokens
+        Output:
+            token_details: list of dicts containing token info and scores
+        """
+        tokens_id = self.tokenizer.encode(text, add_special_tokens=False)
+        if ntoks_max is not None:
+            tokens_id = tokens_id[:ntoks_max]
+        
+        total_len = len(tokens_id)
+        token_details = []
+        seen_grams = set()
+        
+        # Add initial tokens that can't be scored (not enough context)
+        num_start = min(self.ngram, total_len)
+        for i in range(num_start):
+            token_details.append({
+                'token_id': tokens_id[i],
+                'is_scored': False,
+                'score': float('nan'),
+                'token_text': self.tokenizer.decode([tokens_id[i]])
+            })
+        
+        # Score remaining tokens
+        for cur_pos in range(self.ngram, total_len):
+            ngram_tokens = tokens_id[cur_pos-self.ngram:cur_pos]
+            is_scored = True
+            
+            if scoring_method == 'v1':
+                tup_for_unique = tuple(ngram_tokens)
+                is_scored = tup_for_unique not in seen_grams
+                if is_scored:
+                    seen_grams.add(tup_for_unique)
+            elif scoring_method == 'v2':
+                tup_for_unique = tuple(ngram_tokens + [tokens_id[cur_pos]])
+                is_scored = tup_for_unique not in seen_grams
+                if is_scored:
+                    seen_grams.add(tup_for_unique)
+                    
+            score = float('nan')
+            if is_scored:
+                score = self.score_tok(ngram_tokens, tokens_id[cur_pos])
+                score = float(score)
+                
+            token_details.append({
+                'token_id': tokens_id[cur_pos],
+                'is_scored': is_scored,
+                'score': score,
+                'token_text': self.tokenizer.decode([tokens_id[cur_pos]])
+            })
+            
+        return token_details
+
+    def get_pvalues_by_tok(
+        self, 
+        token_details: list[dict]
+    ) -> tuple[list[float], dict]:
+        """
+        Get p-value for each token so far.
+        Args:
+            token_details: list of dicts containing token info and scores from get_details()
+        Returns:
+            tuple containing:
+            - list of p-values, with nan for unscored tokens
+            - dict with auxiliary information:
+                - final_score: final running score
+                - ntoks_scored: final number of scored tokens
+                - final_pvalue: last non-nan pvalue (0.5 if none available)
+        """
+        pvalues = []
+        running_score = 0
+        ntoks_scored = 0
+        eps = 1e-10  # small constant to avoid numerical issues
+        last_valid_pvalue = 0.5  # default value if no tokens are scored
+        
+        for token in token_details:
+            if token['is_scored']:
+                running_score += token['score']
+                ntoks_scored += 1
+                pvalue = self.get_pvalue(running_score, ntoks_scored, eps)
+                last_valid_pvalue = pvalue
+                pvalues.append(pvalue)
+            else:
+                pvalues.append(float('nan'))
+        
+        aux_info = {
+            'final_score': running_score,
+            'ntoks_scored': ntoks_scored,
+            'final_pvalue': last_valid_pvalue
+        }
+            
+        return pvalues, aux_info
+
+    def score_tok(self, ngram_tokens: list[int], token_id: int):
+        """ for each token in the text, compute the score increment """
+        raise NotImplementedError
+    
+    def get_pvalue(self, score: float, ntoks: int, eps: float):
+        """ compute the p-value for a couple of score and number of tokens """
+        raise NotImplementedError
+
+
+class MarylandDetector(WmDetector):
+
+    def __init__(self, 
+            tokenizer: AutoTokenizer,
+            ngram: int = 1,
+            seed: int = 0,
+            gamma: float = 0.5, 
+            delta: float = 1.0, 
+            **kwargs):
+        super().__init__(tokenizer, ngram, seed, **kwargs)
+        self.gamma = gamma
+        self.delta = delta
+    
+    def score_tok(self, ngram_tokens, token_id):
+        """ 
+        score_t = 1 if token_id in greenlist else 0
+        """
+        seed = get_seed_rng(self.seed, ngram_tokens)
+        self.rng.manual_seed(seed)
+        scores = torch.zeros(self.vocab_size)
+        vocab_permutation = torch.randperm(self.vocab_size, generator=self.rng)
+        greenlist = vocab_permutation[:int(self.gamma * self.vocab_size)] # gamma * n toks in the greenlist
+        scores[greenlist] = 1 
+        return scores[token_id]
+                
+    def get_pvalue(self, score: int, ntoks: int, eps: float):
+        """ from cdf of a binomial distribution """
+        pvalue = special.betainc(score, 1 + ntoks - score, self.gamma)
+        return max(pvalue, eps)
+
+class MarylandDetectorZ(WmDetector):
+
+    def __init__(self, 
+            tokenizer: AutoTokenizer,
+            ngram: int = 1,
+            seed: int = 0,
+            gamma: float = 0.5, 
+            delta: float = 1.0, 
+            **kwargs):
+        super().__init__(tokenizer, ngram, seed, **kwargs)
+        self.gamma = gamma
+        self.delta = delta
+    
+    def score_tok(self, ngram_tokens, token_id):
+        """ same as MarylandDetector but using zscore """
+        seed = get_seed_rng(self.seed, ngram_tokens)
+        self.rng.manual_seed(seed)
+        scores = torch.zeros(self.vocab_size)
+        vocab_permutation = torch.randperm(self.vocab_size, generator=self.rng)
+        greenlist = vocab_permutation[:int(self.gamma * self.vocab_size)] # gamma * n
+        scores[greenlist] = 1
+        return scores[token_id]
+                
+    def get_pvalue(self, score: int, ntoks: int, eps: float):
+        """ from cdf of a normal distribution """
+        zscore = (score - self.gamma * ntoks) / np.sqrt(self.gamma * (1 - self.gamma) * ntoks)
+        pvalue = 0.5 * special.erfc(zscore / np.sqrt(2))
+        return max(pvalue, eps)
+    
+class OpenaiDetector(WmDetector):
+
+    def __init__(self, 
+            tokenizer: AutoTokenizer, 
+            ngram: int = 1,
+            seed: int = 0,
+            **kwargs):
+        super().__init__(tokenizer, ngram, seed, **kwargs)
+    
+    def score_tok(self, ngram_tokens, token_id):
+        """ 
+        score_t = -log(1 - rt[token_id]])
+        """
+        seed = get_seed_rng(self.seed, ngram_tokens)
+        self.rng.manual_seed(seed)
+        rs = torch.rand(self.vocab_size, generator=self.rng) # n
+        scores = -(1 - rs).log()
+        return scores[token_id]
+ 
+    def get_pvalue(self, score: float, ntoks: int, eps: float):
+        """ from cdf of a gamma distribution """
+        pvalue = special.gammaincc(ntoks, score)
+        return max(pvalue, eps)
+
+class OpenaiDetectorZ(WmDetector):
+
+    def __init__(self, 
+            tokenizer: AutoTokenizer, 
+            ngram: int = 1,
+            seed: int = 0,
+            **kwargs):
+        super().__init__(tokenizer, ngram, seed, **kwargs)
+    
+    def score_tok(self, ngram_tokens, token_id):
+        """ same as OpenaiDetector but using zscore """
+        seed = get_seed_rng(self.seed, ngram_tokens)
+        self.rng.manual_seed(seed)
+        rs = torch.rand(self.vocab_size, generator=self.rng) # n
+        scores = -(1 - rs).log()
+        return scores[token_id]
+ 
+    def get_pvalue(self, score: float, ntoks: int, eps: float):
+        """ from cdf of a normal distribution """
+        mu0 = 1
+        sigma0 = np.pi / np.sqrt(6)
+        zscore = (score/ntoks - mu0) / (sigma0 / np.sqrt(ntoks))
+        pvalue = 0.5 * special.erfc(zscore / np.sqrt(2))
+        return max(pvalue, eps)

wm_interactive/core/generator.py

@@ -0,0 +1,211 @@
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from .hashing import get_seed_rng
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+class WmGenerator():
+    def __init__(self, 
+        model: AutoModelForCausalLM, 
+        tokenizer: AutoTokenizer, 
+        ngram: int = 1,
+        seed: int = 0,
+        **kwargs
+    ):
+        # model config
+        self.tokenizer = tokenizer
+        self.vocab_size = self.tokenizer.vocab_size
+        self.model = model
+        self.max_seq_len = model.config.max_sequence_length if 'max_sequence_length' in model.config.to_dict() else 2048
+        self.pad_id = model.config.pad_token_id if model.config.pad_token_id is not None else -1
+        self.eos_id = model.config.eos_token_id
+        # watermark config
+        self.ngram = ngram
+        self.seed = seed
+        self.rng = torch.Generator()
+        self.rng.manual_seed(self.seed)
+
+    @torch.no_grad()
+    def generate(
+        self,
+        prompt: str,
+        max_gen_len: int,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        return_aux: bool = False,
+    ) -> str:
+        
+        prompt_tokens = self.tokenizer.encode(prompt)
+        prompt_size = len(prompt_tokens)
+        total_len = min(self.max_seq_len, max_gen_len + prompt_size)
+        tokens = torch.full((1, total_len), self.pad_id).to(device).long()
+        if total_len < prompt_size:
+            print("prompt is bigger than max sequence length")
+            prompt_tokens = prompt_tokens[:total_len]
+        tokens[0, :len(prompt_tokens)] = torch.tensor(prompt_tokens).long()
+        input_text_mask = tokens != self.pad_id
+
+        start_pos = prompt_size
+        prev_pos = 0
+        for cur_pos in range(start_pos, total_len):
+            past_key_values = outputs.past_key_values if prev_pos > 0 else None
+            outputs = self.model.forward(
+                tokens[:, prev_pos:cur_pos], 
+                use_cache=True, 
+                past_key_values=past_key_values
+            )
+            ngram_tokens = tokens[0, cur_pos-self.ngram:cur_pos].tolist()
+            aux = {
+                'ngram_tokens': ngram_tokens,
+                'cur_pos': cur_pos,
+            }
+            next_tok = self.sample_next(outputs.logits[:, -1, :], aux, temperature, top_p)
+            tokens[0, cur_pos] = torch.where(input_text_mask[0, cur_pos], tokens[0, cur_pos], next_tok)
+            prev_pos = cur_pos
+            if next_tok == self.eos_id:
+                break
+
+        # cut to max gen len
+        t = tokens[0, :prompt_size + max_gen_len].tolist()
+        # cut to eos tok if any
+        finish_reason = 'length'
+        try:
+            find_eos = t[prompt_size:].index(self.eos_id)
+            if find_eos:
+                t = t[: prompt_size+find_eos]
+            finish_reason = 'eos'
+        except ValueError:
+            pass
+        aux_info = {
+            't': t, 
+            'finish_reason': finish_reason,
+            'n_toks_gen': len(t) - prompt_size,
+            'n_toks_tot': len(t),
+        }
+        decoded = self.tokenizer.decode(t)
+
+        if return_aux:
+            return decoded, aux_info
+        return decoded
+    
+    def sample_next(
+        self,
+        logits: torch.FloatTensor, # (1, vocab_size): logits for last token
+        aux: dict, # ngram_tokens (1, ngram): tokens to consider when seeding
+        temperature: float = 0.8, # temperature for sampling
+        top_p: float = 0.95, # top p for sampling
+    ):
+        """Vanilla sampling with temperature and top p."""
+        if temperature > 0:
+            probs = torch.softmax(logits / temperature, dim=-1)
+            probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+            probs_sum = torch.cumsum(probs_sort, dim=-1)
+            mask = probs_sum - probs_sort > top_p
+            probs_sort[mask] = 0.0
+            probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+            next_token = torch.multinomial(probs_sort, num_samples=1) # one hot of next token, ordered by original probs
+            next_token = torch.gather(probs_idx, -1, next_token) # one hot of next token, ordered by vocab
+        else:
+            next_token = torch.argmax(logits, dim=-1)
+        next_token = next_token.reshape(-1)[0]  # Get the single token value
+        return next_token
+
+
+class OpenaiGenerator(WmGenerator):
+    """
+    Generate text using LLaMA and Aaronson's watermarking method.
+    From ngram tokens, select the next token based on the following:
+    - hash the ngram tokens and get a seed
+    - use the seed to generate V random number r between [0,1]
+    - select argmax ( r^(1/p) )
+    """
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)        
+
+    def sample_next(
+        self,
+        logits: torch.FloatTensor, # (1, vocab_size): logits for last token
+        aux: dict, # (1, ngram): tokens to consider when seeding
+        temperature: float = 0.8, # temperature for sampling
+        top_p: float = 0.95, # top p for sampling
+    ):
+        ngram_tokens = aux['ngram_tokens']
+        if temperature > 0:
+            probs = torch.softmax(logits / temperature, dim=-1)
+            probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+            probs_sum = torch.cumsum(probs_sort, dim=-1)
+            mask = probs_sum - probs_sort > top_p
+            probs_sort[mask] = 0.0
+            probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+            # seed with hash of ngram tokens
+            seed = get_seed_rng(self.seed, ngram_tokens)
+            self.rng.manual_seed(seed)
+            # generate rs randomly between [0,1]
+            rs = torch.rand(self.vocab_size, generator=self.rng) # n
+            rs = torch.Tensor(rs).to(probs_sort.device)
+            rs = rs[probs_idx[0]] 
+            # compute r^(1/p)
+            probs_sort[0] = torch.pow(rs, 1/probs_sort[0])
+            # select argmax ( r^(1/p) )
+            next_token = torch.argmax(probs_sort, dim=-1, keepdim=True)
+            next_token = torch.gather(probs_idx, -1, next_token)
+        else:
+            next_token = torch.argmax(logits, dim=-1)
+        next_token = next_token.reshape(-1)[0]  # Get the single token value
+        return next_token
+
+
+class MarylandGenerator(WmGenerator):
+    """
+    Generate text using LLaMA and Maryland's watemrarking method.
+    From ngram tokens, select the next token based on the following:
+    - hash the ngram tokens and get a seed
+    - use the seed to partition the vocabulary into greenlist (gamma*V words) and blacklist 
+    - add delta to greenlist words' logits
+    """
+    def __init__(self, 
+            *args, 
+            gamma: float = 0.5,
+            delta: float = 1.0,
+            **kwargs
+        ):
+        super().__init__(*args, **kwargs)        
+        self.gamma = gamma
+        self.delta = delta
+
+    def sample_next(
+        self,
+        logits: torch.FloatTensor, # (1, vocab_size): logits for last token
+        aux: dict, # ngram_tokens (1, ngram): tokens to consider when seeding
+        temperature: float = 0.8, # temperature for sampling
+        top_p: float = 0.95, # top p for sampling
+    ):
+        ngram_tokens = aux['ngram_tokens']
+        logits = self.logits_processor(logits, ngram_tokens)
+        if temperature > 0:
+            probs = torch.softmax(logits / temperature, dim=-1)
+            probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)
+            probs_sum = torch.cumsum(probs_sort, dim=-1)
+            mask = probs_sum - probs_sort > top_p
+            probs_sort[mask] = 0.0
+            probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))
+            next_token = torch.multinomial(probs_sort, num_samples=1) # one hot of next token, ordered by original probs
+            next_token = torch.gather(probs_idx, -1, next_token) # one hot of next token, ordered by vocab
+        else:
+            next_token = torch.argmax(logits, dim=-1)
+        next_token = next_token.reshape(-1)[0]  # Get the single token value
+        return next_token
+
+    def logits_processor(self, logits, ngram_tokens):
+        """Process logits to mask out words in greenlist."""
+        logits = logits.clone()
+        seed = get_seed_rng(self.seed, ngram_tokens)
+        self.rng.manual_seed(seed)
+        vocab_permutation = torch.randperm(self.vocab_size, generator=self.rng)
+        greenlist = vocab_permutation[:int(self.gamma * self.vocab_size)] # gamma * n
+        bias = torch.zeros(self.vocab_size).to(logits.device)
+        bias[greenlist] = self.delta
+        logits[0] += bias # add bias to greenlist words
+        return logits

wm_interactive/core/hashing.py

@@ -0,0 +1,13 @@
+
+def get_seed_rng(
+    start, 
+    input_ids: list[int],
+    salt = 35317
+) -> int:
+    """
+    Seed RNG with hash of input_ids.
+    Adapted from https://github.com/jwkirchenbauer/lm-watermarking
+    """
+    for ii in input_ids:
+        start = (start * salt + ii) % (2 ** 64 - 1)
+    return int(start)

wm_interactive/core/main.py

@@ -0,0 +1,256 @@
+"""
+Main script for watermark detection.
+Test with:
+    python -m wm_interactive.core.main --model_name smollm2-135m --prompt_path data/prompts.json --method maryland --delta 4.0 --ngram 1
+"""
+
+import os
+import json
+import time
+import tqdm
+import torch
+import numpy as np
+import pandas as pd
+import argparse
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from wm_interactive.core.generator import WmGenerator, OpenaiGenerator, MarylandGenerator
+from wm_interactive.core.detector import WmDetector, OpenaiDetector, OpenaiDetectorZ, MarylandDetector, MarylandDetectorZ
+
+# model names mapping
+model_names = {
+    # 'llama-3.2-1b': 'meta-llama/Llama-3.2-1B-Instruct',
+    'smollm2-135m': 'HuggingFaceTB/SmolLM2-135M-Instruct',
+    'smollm2-360m': 'HuggingFaceTB/SmolLM2-360M-Instruct',   
+}
+
+CACHE_DIR = "wm_interactive/static/hf_cache"
+
+
+def load_prompts(json_path: str, prompt_type: str = "smollm", nsamples: int = None) -> list[dict]:
+    """Load prompts from a JSON file.
+    
+    Args:
+        json_path: Path to the JSON file
+        prompt_type: Type of prompt dataset (alpaca, smollm)
+        nsamples: Number of samples to load (if None, load all)
+    
+    Returns:
+        List of prompts
+    """
+    if not os.path.exists(json_path):
+        raise FileNotFoundError(f"File {json_path} not found")
+    
+    with open(json_path, 'r') as f:
+        data = json.load(f)
+    
+    if prompt_type == "alpaca":
+        prompts = [{"instruction": item["instruction"]} for item in data]
+    elif prompt_type == "smollm":
+        prompts = []
+        for item in data:
+            prompt = "<|im_start|>system\nYou are a helpful AI assistant named SmolLM, trained by Hugging Face<|im_end|>\n"
+            prompt += f"<|im_start|>user\n{item['instruction']}<|im_end|>\n<|im_start|>assistant\n"
+            prompts.append({"instruction": prompt})
+    else:
+        raise ValueError(f"Prompt type {prompt_type} not supported")
+    
+    if nsamples is not None:
+        prompts = prompts[:nsamples]
+    
+    return prompts 
+
+def load_results(json_path: str, result_key: str = "result", nsamples: int = None) -> list[str]:
+    """Load results from a JSONL file.
+    
+    Args:
+        json_path: Path to the JSONL file
+        result_key: Key to extract from each JSON line
+        nsamples: Number of samples to load (if None, load all)
+    
+    Returns:
+        List of results
+    """
+    if not os.path.exists(json_path):
+        raise FileNotFoundError(f"File {json_path} not found")
+    
+    results = []
+    with open(json_path, 'r') as f:
+        for line in f:
+            if line.strip():  # Skip empty lines
+                data = json.loads(line)
+                results.append(data[result_key])
+            if nsamples is not None and len(results) >= nsamples:
+                break
+    
+    return results
+
+def get_args_parser():
+    parser = argparse.ArgumentParser('Args', add_help=False)
+
+    # model parameters
+    parser.add_argument('--model_name', type=str, required=True, 
+                       help='Name of the model to use. Choose from: llama-3.2-1b, smollm2-135m')
+
+    # prompts parameters
+    parser.add_argument('--prompt_path', type=str, default=None,
+                       help='Path to the prompt dataset. Required if --prompt is not provided')
+    parser.add_argument('--prompt_type', type=str, default="smollm",
+                       help='Type of prompt dataset. Only used if --prompt_path is provided')
+    parser.add_argument('--prompt', type=str, nargs='+', default=None,
+                       help='List of prompts to use. If not provided, prompts will be loaded from --prompt_path')
+
+    # generation parameters
+    parser.add_argument('--temperature', type=float, default=0.8,
+                       help='Temperature for sampling (higher = more random)')
+    parser.add_argument('--top_p', type=float, default=0.95,
+                       help='Top p for nucleus sampling (lower = more focused)')
+    parser.add_argument('--max_gen_len', type=int, default=256,
+                       help='Maximum length of generated text')
+
+    # watermark parameters
+    parser.add_argument('--method', type=str, default='none',
+                       help='Watermarking method. Choose from: none (no watermarking), openai (Aaronson et al.), maryland (Kirchenbauer et al.)')
+    parser.add_argument('--method_detect', type=str, default='same',
+                       help='Statistical test to detect watermark. Choose from: same (same as method), openai, openaiz, maryland, marylandz')
+    parser.add_argument('--seed', type=int, default=0,
+                       help='Random seed for reproducibility')
+    parser.add_argument('--ngram', type=int, default=1,
+                       help='n-gram size for rng key generation')
+    parser.add_argument('--gamma', type=float, default=0.5,
+                       help='For maryland method: proportion of greenlist tokens')
+    parser.add_argument('--delta', type=float, default=2.0,
+                       help='For maryland method: bias to add to greenlist tokens')
+    parser.add_argument('--scoring_method', type=str, default='v2',
+                       help='Method for scoring. Choose from: none (score every token), v1 (score when context unique), v2 (score when context+token unique)')
+
+    # experiment parameters
+    parser.add_argument('--nsamples', type=int, default=None,
+                       help='Number of samples to generate from the prompt dataset')
+    parser.add_argument('--do_eval', type=bool, default=True,
+                       help='Whether to evaluate the generated text')
+    parser.add_argument('--output_dir', type=str, default='output',
+                       help='Directory to save results')
+
+    return parser
+
+def main(args):
+    print('job dir: {}'.format(os.path.dirname(os.path.realpath(__file__))))
+    print("{}".format(args).replace(', ', ',\n'))
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    # build model
+    model_name = args.model_name.lower()
+    if model_name not in model_names:
+        raise ValueError(f"Model {model_name} not supported. Choose from: {list(model_names.keys())}")
+    model_name = model_names[model_name]
+    
+    # Load tokenizer and model
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=CACHE_DIR)
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        cache_dir=CACHE_DIR
+    ).to(device)
+
+    # build watermark generator
+    if args.method == "none":
+        generator = WmGenerator(model, tokenizer)
+    elif args.method == "openai":
+        generator = OpenaiGenerator(model, tokenizer, args.ngram, args.seed)
+    elif args.method == "maryland":
+        generator = MarylandGenerator(model, tokenizer, args.ngram, args.seed, gamma=args.gamma, delta=args.delta)
+    else:
+        raise NotImplementedError("method {} not implemented".format(args.method))
+
+    # load prompts
+    if args.prompt is not None:
+        prompts = args.prompt
+        prompts = [{"instruction": prompt} for prompt in prompts]
+    elif args.prompt_path is not None:
+        prompts = load_prompts(json_path=args.prompt_path, prompt_type=args.prompt_type, nsamples=args.nsamples)
+    else:
+        raise ValueError("Either --prompt or --prompt_path must be provided")
+    
+    # (re)start experiment
+    os.makedirs(args.output_dir, exist_ok=True)
+    start_point = 0 # if resuming, start from the last line of the file
+    if os.path.exists(os.path.join(args.output_dir, f"results.jsonl")):
+        with open(os.path.join(args.output_dir, f"results.jsonl"), "r") as f:
+            for _ in f:
+                start_point += 1
+    print(f"Starting from {start_point}")
+
+    # generate
+    all_times = []
+    with open(os.path.join(args.output_dir, f"results.jsonl"), "a") as f:
+        for ii in range(start_point, len(prompts)):
+            # generate text
+            time0 = time.time()
+            prompt = prompts[ii]["instruction"]
+            result = generator.generate(
+                prompt, 
+                max_gen_len=args.max_gen_len, 
+                temperature=args.temperature, 
+                top_p=args.top_p
+            )
+            time1 = time.time()
+            # time chunk
+            speed = 1 / (time1 - time0)
+            eta = (len(prompts) - ii) / speed
+            eta = time.strftime("%Hh%Mm%Ss", time.gmtime(eta)) 
+            all_times.append(time1 - time0)
+            print(f"Generated {ii:5d} - Speed {speed:.2f} prompts/s - ETA {eta}")
+            # log
+            f.write(json.dumps({
+                "prompt": prompt, 
+                "result": result[len(prompt):],
+                "speed": speed,
+                "eta": eta}) + "\n")
+            f.flush()
+    print(f"Average time per prompt: {np.sum(all_times) / (len(prompts) - start_point) :.2f}")
+
+    if args.method_detect == 'same':
+        args.method_detect = args.method
+    if (not args.do_eval) or (args.method_detect not in ["openai", "maryland", "marylandz", "openaiz"]):
+        return
+    
+    # build watermark detector
+    if args.method_detect == "openai":
+        detector = OpenaiDetector(tokenizer, args.ngram, args.seed)
+    elif args.method_detect == "openaiz":
+        detector = OpenaiDetectorZ(tokenizer, args.ngram, args.seed)
+    elif args.method_detect == "maryland":
+        detector = MarylandDetector(tokenizer, args.ngram, args.seed, gamma=args.gamma, delta=args.delta)
+    elif args.method_detect == "marylandz":
+        detector = MarylandDetectorZ(tokenizer, args.ngram, args.seed, gamma=args.gamma, delta=args.delta)
+
+    # evaluate
+    results = load_results(json_path=os.path.join(args.output_dir, f"results.jsonl"), result_key="result", nsamples=args.nsamples)
+    log_stats = []
+    with open(os.path.join(args.output_dir, 'scores.jsonl'), 'w') as f:
+        for text in tqdm.tqdm(results):
+            # get token details and pvalues
+            token_details = detector.get_details(text, scoring_method=args.scoring_method)
+            pvalues, aux_info = detector.get_pvalues_by_tok(token_details)
+            # log stats
+            log_stat = {
+                'num_token': aux_info['ntoks_scored'],
+                'score': aux_info['final_score'],
+                'pvalue': aux_info['final_pvalue'],
+                'log10_pvalue': np.log10(aux_info['final_pvalue']),
+            }
+            log_stats.append(log_stat)
+            f.write('\n' + json.dumps({k: float(v) for k, v in log_stat.items()}))
+        df = pd.DataFrame(log_stats)
+        print(f">>> Scores: \n{df.describe(percentiles=[])}")
+        print(f"Saved scores to {os.path.join(args.output_dir, 'scores.csv')}")
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args)

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/0ad5ecc2035b7031b88afb544ee95e2d49baa484

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82b84012e3add4d01d12ba14442026e49b8cbbaead1f79ecf3d919784f82dc79
+size 800662

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/36293b6099200eb8aeb55ae2c01bca2ba46d80d0

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8eb740e8bbe4cff95ea7b4588d17a2432deb16e8075bc5828ff7ba9be94d982a
+size 861

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/44719d2e365acac0637fd25a3acf46494ca45940

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b7379f3ae813529281a5c602bc5a11c1d4e0a99107aaa597fe936c1e813ca52
+size 655

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/5af571cbf074e6d21a03528d2330792e532ca608f24ac70a143f6b369968ab8c

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5af571cbf074e6d21a03528d2330792e532ca608f24ac70a143f6b369968ab8c
+size 269060552

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/69503b13f727ba3812b6803e97442a6de05ef5eb

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b54e8aa4e53d5383e2e4bc635a56b43f9647f7b13832d5d9ecd8f82dac4f510
+size 466391

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/8c7b22013909450429303ed10be4398bd63f5457

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ec77d44f62efeb38d7e044a1db318f6a939438425312dfa333b8382dbad98df
+size 3764

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/da6c4d71a43aa7e6f785bdbb28ea5025438a73fa

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87b916edaaab66b3899b9d0dd0752727dff6666686da0504d89ae0a6e055a013
+size 132

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/blobs/f922b1797f0c88e71addc8393787831f2477a4bd

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ca9acddb6525a194ec8ac7a87f24fbba7232a9a15ffa1af0c1224fcd888e47c
+size 2104556

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/refs/main

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71a184f20b0fe5c1a9407ed75fa9633b681779c7f1a5ca478f22fdff69a6c7ab
+size 40

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8eb740e8bbe4cff95ea7b4588d17a2432deb16e8075bc5828ff7ba9be94d982a
+size 861

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/generation_config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:87b916edaaab66b3899b9d0dd0752727dff6666686da0504d89ae0a6e055a013
+size 132

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/merges.txt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0b54e8aa4e53d5383e2e4bc635a56b43f9647f7b13832d5d9ecd8f82dac4f510
+size 466391

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5af571cbf074e6d21a03528d2330792e532ca608f24ac70a143f6b369968ab8c
+size 269060552

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/special_tokens_map.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b7379f3ae813529281a5c602bc5a11c1d4e0a99107aaa597fe936c1e813ca52
+size 655

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9ca9acddb6525a194ec8ac7a87f24fbba7232a9a15ffa1af0c1224fcd888e47c
+size 2104556

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/tokenizer_config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ec77d44f62efeb38d7e044a1db318f6a939438425312dfa333b8382dbad98df
+size 3764

wm_interactive/static/hf_cache/models--HuggingFaceTB--SmolLM2-135M-Instruct/snapshots/e2c3f7557efbdec707ae3a336371d169783f1da1/vocab.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:82b84012e3add4d01d12ba14442026e49b8cbbaead1f79ecf3d919784f82dc79
+size 800662

wm_interactive/static/styles.css

@@ -0,0 +1,357 @@
+body {
+    background-color: #f7f7f8;
+    color: #1a1a1a;
+    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
+    line-height: 1.5;
+    padding: 0;
+    margin: 0;
+}
+
+.container {
+    background-color: transparent;
+    box-shadow: none;
+    max-width: 1000px;
+    padding: 20px;
+    margin: 0 auto;
+}
+
+h1 {
+    color: #1a1a1a;
+    font-size: 24px;
+    font-weight: 600;
+    margin-bottom: 30px;
+}
+
+.input-section {
+    display: flex;
+    flex-direction: column;
+    gap: 24px;
+    margin-bottom: 30px;
+}
+
+.input-section textarea {
+    width: 100%;
+    padding: 16px;
+    background-color: #ffffff;
+    border: 1px solid #e5e5e5;
+    border-radius: 12px;
+    resize: none;
+    font-size: 16px;
+    line-height: 1.5;
+    color: #1a1a1a;
+    transition: border-color 0.2s;
+}
+
+.input-section textarea:focus {
+    outline: none;
+    border-color: #10a37f;
+    box-shadow: 0 0 0 2px rgba(16, 163, 127, 0.2);
+}
+
+.input-section #prompt_text {
+    height: 80px;
+    padding-right: 52px;
+}
+
+.input-section #user_text {
+    height: 160px;
+}
+
+.button-container {
+    display: flex;
+    gap: 12px;
+    justify-content: center;
+}
+
+.btn {
+    padding: 8px 16px;
+    font-size: 14px;
+    font-weight: 500;
+    border-radius: 6px;
+    transition: all 0.2s;
+}
+
+.btn-primary {
+    background-color: #10a37f;
+    border-color: #10a37f;
+}
+
+.btn-primary:hover:not(:disabled) {
+    background-color: #0e8d6e;
+    border-color: #0e8d6e;
+}
+
+.btn-secondary {
+    background-color: #40414f;
+    border-color: #565869;
+    color: #ececf1;
+}
+
+.btn-secondary:hover:not(:disabled) {
+    background-color: #4a4b5a;
+    border-color: #6b6c7b;
+}
+
+.token-display {
+    margin: 24px 0;
+    padding: 16px;
+    background-color: #ffffff;
+    border: 1px solid #e5e5e5;
+    border-radius: 12px;
+    min-height: 100px;
+    font-size: 15px;
+    line-height: 1.6;
+}
+
+.stats-container {
+    display: grid;
+    grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+    gap: 20px;
+    margin-top: 30px;
+    padding: 20px;
+    background-color: #ffffff;
+    border-radius: 12px;
+    border: 1px solid #e5e5e5;
+}
+
+.stats-container > div {
+    text-align: center;
+    padding: 16px;
+    border-radius: 8px;
+    background-color: #f7f7f8;
+}
+
+.stat-value {
+    font-size: 28px;
+    font-weight: 600;
+    color: #1a1a1a;
+    margin-bottom: 8px;
+}
+
+.stat-label {
+    position: relative;
+    color: #6e6e80;
+    font-size: 14px;
+    display: inline-flex;
+    align-items: center;
+    justify-content: center;
+    gap: 6px;
+}
+
+.help-icon {
+    color: #6e6e80;
+    font-size: 12px;
+    opacity: 0.8;
+    transition: opacity 0.2s;
+    cursor: help;
+}
+
+.help-tooltip {
+    visibility: hidden;
+    position: absolute;
+    z-index: 1000;
+    bottom: 125%;
+    left: 50%;
+    transform: translateX(-50%);
+    background-color: #1a1a1a;
+    color: #ffffff;
+    padding: 8px 12px;
+    border-radius: 6px;
+    font-size: 12px;
+    width: max-content;
+    max-width: 200px;
+    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3);
+    pointer-events: none;
+    opacity: 0;
+    transition: opacity 0.2s;
+}
+
+.help-tooltip::after {
+    content: "";
+    position: absolute;
+    top: 100%;
+    left: 50%;
+    margin-left: -5px;
+    border-width: 5px;
+    border-style: solid;
+    border-color: #1a1a1a transparent transparent transparent;
+}
+
+.help-icon:hover + .help-tooltip {
+    visibility: visible;
+    opacity: 1;
+}
+
+.token {
+    padding: 2px 4px;
+    margin: 1px;
+    border-radius: 4px;
+    font-family: 'SF Mono', 'Menlo', 'Monaco', Courier, monospace;
+    transition: background-color 0.2s;
+    position: relative;
+    cursor: pointer;
+}
+
+.token:hover {
+    filter: brightness(1.1);
+}
+
+.token-tooltip {
+    visibility: hidden;
+    position: absolute;
+    z-index: 1000;
+    bottom: 125%;
+    left: 50%;
+    transform: translateX(-50%);
+    background-color: #1a1a1a;
+    color: #ffffff;
+    padding: 8px 12px;
+    border-radius: 6px;
+    font-size: 12px;
+    white-space: nowrap;
+    box-shadow: 0 4px 12px rgba(0, 0, 0, 0.3);
+    pointer-events: none;
+}
+
+.token-tooltip::after {
+    content: "";
+    position: absolute;
+    top: 100%;
+    left: 50%;
+    margin-left: -5px;
+    border-width: 5px;
+    border-style: solid;
+    border-color: #1a1a1a transparent transparent transparent;
+}
+
+.token:hover .token-tooltip {
+    visibility: visible;
+}
+
+/* Modal styling */
+.modal-content {
+    background-color: #ffffff;
+    color: #1a1a1a;
+    border: 1px solid #e5e5e5;
+}
+
+.modal-header {
+    border-bottom-color: #e5e5e5;
+}
+
+.modal-footer {
+    border-top-color: #e5e5e5;
+}
+
+.form-control, .form-select {
+    background-color: #f7f7f8;
+    border-color: #e5e5e5;
+    color: #1a1a1a;
+}
+
+.form-control:focus, .form-select:focus {
+    background-color: #f7f7f8;
+    border-color: #10a37f;
+    color: #1a1a1a;
+    box-shadow: 0 0 0 2px rgba(16, 163, 127, 0.2);
+}
+
+.form-text {
+    color: #6e6e80;
+}
+
+.btn-close {
+    filter: none;
+}
+
+/* Mobile-specific styles */
+@media (max-width: 768px) {
+    .container {
+        padding: 15px;
+    }
+
+    .stats-container {
+        grid-template-columns: repeat(2, 1fr);
+        gap: 12px;
+        padding: 12px;
+    }
+
+    .stat-value {
+        font-size: 24px;
+    }
+
+    .stat-label {
+        font-size: 12px;
+    }
+
+    .help-tooltip {
+        max-width: 160px;
+    }
+}
+
+/* Light scrollbar */
+::-webkit-scrollbar {
+    width: 8px;
+    height: 8px;
+}
+
+::-webkit-scrollbar-track {
+    background: #f7f7f8;
+}
+
+::-webkit-scrollbar-thumb {
+    background: #d1d1d1;
+    border-radius: 4px;
+}
+
+::-webkit-scrollbar-thumb:hover {
+    background: #a8a8a8;
+}
+
+.prompt-container {
+    position: relative;
+    width: 100%;
+}
+
+.floating-btn {
+    position: absolute;
+    bottom: 16px;
+    right: 16px;
+    width: 36px;
+    height: 36px;
+    border-radius: 12px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    border: none;
+    background-color: #10a37f;
+    color: #ffffff;
+    cursor: pointer;
+    transition: all 0.2s ease;
+    padding: 0;
+}
+
+.floating-btn:hover {
+    background-color: #0e8d6e;
+}
+
+.floating-btn:disabled {
+    background-color: #565869;
+    cursor: not-allowed;
+}
+
+.floating-btn i {
+    font-size: 16px;
+}
+
+.floating-btn .stop-icon {
+    display: none;
+}
+
+.floating-btn.generating .send-icon {
+    display: none;
+}
+
+.floating-btn.generating .stop-icon {
+    display: block;
+}

wm_interactive/templates/index.html

@@ -0,0 +1,459 @@
+<!DOCTYPE html>
+<html>
+<head>
+    <title>Watermark Detector</title>
+    <meta name="viewport" content="width=device-width, initial-scale=1">
+    <link href="https://cdn.jsdelivr.net/npm/[email protected]/dist/css/bootstrap.min.css" rel="stylesheet">
+    <link rel="stylesheet" href="https://cdn.jsdelivr.net/npm/[email protected]/font/bootstrap-icons.css">
+    <link rel="stylesheet" href="{{ url_for('static', filename='styles.css') }}">
+</head>
+<body>
+    <div class="container">
+        <div class="d-flex justify-content-between align-items-center">
+            <h1>Interactive watermark detector</h1>
+            <button class="btn btn-outline-secondary" data-bs-toggle="modal" data-bs-target="#paramsModal">
+                <i class="bi bi-gear"></i>
+            </button>
+        </div>
+
+        <!-- Advanced Parameters Modal -->
+        <div class="modal fade" id="paramsModal" tabindex="-1">
+            <div class="modal-dialog">
+                <div class="modal-content">
+                    <div class="modal-header">
+                        <h5 class="modal-title">Advanced Parameters</h5>
+                        <button type="button" class="btn-close" data-bs-dismiss="modal"></button>
+                    </div>
+                    <div class="modal-body">
+                        <div class="mb-3">
+                            <label for="detectorType" class="form-label">Detector Type</label>
+                            <select class="form-select" id="detectorType">
+                                <option value="maryland">Maryland</option>
+                                <option value="marylandz">Maryland Z-score</option>
+                                <option value="openai">OpenAI</option>
+                                <option value="openaiz">OpenAI Z-score</option>
+                            </select>
+                            <div class="form-text">Type of watermark detection algorithm</div>
+                        </div>
+                        <div class="mb-3">
+                            <label for="seed" class="form-label">Seed</label>
+                            <input type="number" class="form-control" id="seed" value="0">
+                            <div class="form-text">Random seed for the watermark detector</div>
+                        </div>
+                        <div class="mb-3">
+                            <label for="ngram" class="form-label">N-gram Size</label>
+                            <input type="number" class="form-control" id="ngram" value="1">
+                            <div class="form-text">Size of the n-gram window used for detection</div>
+                        </div>
+                        <div class="mb-3">
+                            <label for="delta" class="form-label">Delta</label>
+                            <input type="number" step="0.1" class="form-control" id="delta" value="2.0">
+                            <div class="form-text">Bias added to greenlist tokens (for Maryland method)</div>
+                        </div>
+                        <div class="mb-3">
+                            <label for="temperature" class="form-label">Temperature</label>
+                            <input type="number" step="0.1" class="form-control" id="temperature" value="0.8">
+                            <div class="form-text">Temperature for sampling (higher = more random)</div>
+                        </div>
+                    </div>
+                    <div class="modal-footer">
+                        <button type="button" class="btn btn-secondary" data-bs-dismiss="modal">Close</button>
+                        <button type="button" class="btn btn-primary" id="applyParams">Apply</button>
+                    </div>
+                </div>
+            </div>
+        </div>
+
+        <!-- Input Form -->
+        <div class="input-section">
+            <div class="prompt-container">
+                <textarea id="prompt_text" 
+                    placeholder="Enter your prompt here to generate text with the model..."></textarea>
+                <button class="floating-btn" id="generateBtn">
+                    <i class="bi bi-send-fill send-icon"></i>
+                    <i class="bi bi-stop-fill stop-icon"></i>
+                </button>
+            </div>
+            <textarea id="user_text" 
+                placeholder="Generated text will appear here. Replace or edit this text to see how watermark detection works."></textarea>
+        </div>
+
+        <!-- Token Display -->
+        <div class="token-display" id="tokenDisplay"></div>
+
+        <!-- Statistics -->
+        <div class="stats-container">
+            <div>
+                <div class="stat-value" id="tokenCount">0</div>
+                <div class="stat-label">
+                    Tokens
+                    <i class="bi bi-question-circle help-icon"></i>
+                    <span class="help-tooltip">Total number of tokens in the text</span>
+                </div>
+            </div>
+            <div>
+                <div class="stat-value" id="scoredTokens">0</div>
+                <div class="stat-label">
+                    Scored Tokens
+                    <i class="bi bi-question-circle help-icon"></i>
+                    <span class="help-tooltip">Number of tokens that were actually scored by the detector (excludes first n-gram tokens and duplicates)</span>
+                </div>
+            </div>
+            <div>
+                <div class="stat-value" id="finalScore">0.00</div>
+                <div class="stat-label">
+                    Final Score
+                    <i class="bi bi-question-circle help-icon"></i>
+                    <span class="help-tooltip">Cumulative score from all scored tokens. Higher values indicate more likely watermarked text</span>
+                </div>
+            </div>
+            <div>
+                <div class="stat-value" id="pValue">0.500</div>
+                <div class="stat-label">
+                    P-value
+                    <i class="bi bi-question-circle help-icon"></i>
+                    <span class="help-tooltip">Statistical significance of the score. Lower values indicate stronger evidence of watermarking (p < 0.05 is typically considered significant)</span>
+                </div>
+            </div>
+        </div>
+    </div>
+
+    <script src="https://cdn.jsdelivr.net/npm/[email protected]/dist/js/bootstrap.bundle.min.js"></script>
+    <script>
+        let debounceTimeout = null;
+        let abortController = null;
+        const textarea = document.getElementById('user_text');
+        const promptArea = document.getElementById('prompt_text');
+        const generateBtn = document.getElementById('generateBtn');
+        const tokenDisplay = document.getElementById('tokenDisplay');
+        const tokenCount = document.getElementById('tokenCount');
+        const scoredTokens = document.getElementById('scoredTokens');
+        const finalScore = document.getElementById('finalScore');
+        const pValue = document.getElementById('pValue');
+        const applyParamsBtn = document.getElementById('applyParams');
+        const seedInput = document.getElementById('seed');
+        const ngramInput = document.getElementById('ngram');
+        const detectorTypeSelect = document.getElementById('detectorType');
+        const deltaInput = document.getElementById('delta');
+        const temperatureInput = document.getElementById('temperature');
+
+        function startGeneration() {
+            const prompt = promptArea.value.trim();
+            if (!prompt) {
+                alert('Please enter a prompt first.');
+                return;
+            }
+
+            generateBtn.classList.add('generating');
+            textarea.value = '';
+
+            // Create new AbortController for this request
+            abortController = new AbortController();
+
+            // Get current parameters
+            const params = {
+                detector_type: detectorTypeSelect.value,
+                seed: parseInt(seedInput.value) || 0,
+                ngram: parseInt(ngramInput.value) || 1,
+                delta: parseFloat(deltaInput.value) || 2.0,
+                temperature: parseFloat(temperatureInput.value) || 0.8
+            };
+
+            // Create headers for SSE
+            const headers = new Headers({
+                'Content-Type': 'application/json',
+                'Accept': 'text/event-stream',
+            });
+
+            // Start fetch request with signal
+            fetch('/generate', {
+                method: 'POST',
+                headers: headers,
+                body: JSON.stringify({ 
+                    prompt: prompt,
+                    params: params
+                }),
+                signal: abortController.signal  // Add the abort signal
+            }).then(response => {
+                const reader = response.body.getReader();
+                const decoder = new TextDecoder();
+                let buffer = '';
+
+                function processText(text) {
+                    const lines = text.split('\n');
+                    
+                    for (const line of lines) {
+                        if (line.startsWith('data: ')) {
+                            try {
+                                const data = JSON.parse(line.slice(6));
+                                
+                                if (data.error) {
+                                    alert('Error: ' + data.error);
+                                    stopGeneration();
+                                    return;
+                                }
+
+                                if (data.token) {
+                                    // Append new token to existing text
+                                    textarea.value += data.token;
+                                    updateTokenization();
+                                }
+
+                                if (data.text) {
+                                    // Final text (only used if something went wrong with streaming)
+                                    textarea.value = data.text;
+                                    updateTokenization();
+                                }
+
+                                if (data.done) {
+                                    stopGeneration();
+                                }
+                            } catch (e) {
+                                console.error('Error parsing SSE data:', e);
+                            }
+                        }
+                    }
+                }
+
+                function pump() {
+                    return reader.read().then(({value, done}) => {
+                        if (done) {
+                            if (buffer.length > 0) {
+                                processText(buffer);
+                            }
+                            return;
+                        }
+
+                        buffer += decoder.decode(value, {stream: true});
+                        const lines = buffer.split('\n\n');
+                        buffer = lines.pop();
+                        
+                        for (const line of lines) {
+                            processText(line);
+                        }
+
+                        return pump();
+                    });
+                }
+
+                return pump();
+            })
+            .catch(error => {
+                if (error.name === 'AbortError') {
+                    console.log('Generation stopped by user');
+                } else {
+                    console.error('Error:', error);
+                    alert('Error: Failed to generate text');
+                }
+            })
+            .finally(() => {
+                generateBtn.classList.remove('generating');
+                abortController = null;
+            });
+        }
+
+        function stopGeneration() {
+            if (abortController) {
+                abortController.abort();
+                abortController = null;
+            }
+            generateBtn.classList.remove('generating');
+        }
+
+        // Remove BOTH old event listeners and add just one new one
+        generateBtn.addEventListener('click', function(e) {
+            e.preventDefault(); // Prevent any double triggers
+            if (generateBtn.classList.contains('generating')) {
+                stopGeneration();
+            } else {
+                startGeneration();
+            }
+        });
+
+        async function updateTokenization() {
+            const text = textarea.value;
+            try {
+                // Validate parameters before sending
+                const seed = parseInt(seedInput.value);
+                const ngram = parseInt(ngramInput.value);
+                const delta = parseFloat(deltaInput.value);
+                const temperature = parseFloat(temperatureInput.value);
+                
+                const response = await fetch('/tokenize', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({ 
+                        text: text,
+                        params: {
+                            detector_type: detectorTypeSelect.value,
+                            seed: isNaN(seed) ? 0 : seed,
+                            ngram: isNaN(ngram) ? 1 : ngram,
+                            delta: isNaN(delta) ? 2.0 : delta,
+                            temperature: isNaN(temperature) ? 0.8 : temperature
+                        }
+                    })
+                });
+
+                if (!response.ok) {
+                    const errorData = await response.json();
+                    throw new Error(errorData.error || `HTTP error! status: ${response.status}`);
+                }
+                
+                const data = await response.json();
+                
+                if (data.error) {
+                    throw new Error(data.error);
+                }
+                
+                // Update token display
+                tokenDisplay.innerHTML = data.tokens.map((token, i) => {
+                    const score = data.scores[i];
+                    const pvalue = data.pvalues[i];
+                    const scoreDisplay = (score !== null && !isNaN(score)) ? score.toFixed(3) : 'N/A';
+                    const pvalueDisplay = (pvalue !== null && !isNaN(pvalue)) ? formatPValue(pvalue) : 'N/A';
+                    
+                    return `<span class="token" style="background-color: ${data.colors[i]}">
+                        ${token}
+                        <div class="token-tooltip">
+                            Score: ${scoreDisplay}<br>
+                            P-value: ${pvalueDisplay}
+                        </div>
+                    </span>`;
+                }).join('');
+                
+                // Update counts and stats - safely handle null values
+                tokenCount.textContent = data.token_count || 0;
+                scoredTokens.textContent = data.ntoks_scored || 0;
+                finalScore.textContent = (data.final_score !== null && !isNaN(data.final_score)) ? 
+                    data.final_score.toFixed(2) : '0.00';
+                pValue.textContent = (data.final_pvalue !== null && !isNaN(data.final_pvalue)) ? 
+                    formatPValue(data.final_pvalue) : '0.500';
+
+                // Clear any previous error
+                const existingError = tokenDisplay.querySelector('.alert-danger');
+                if (existingError) {
+                    existingError.remove();
+                }
+            } catch (error) {
+                console.error('Error updating tokenization:', error);
+                // Show detailed error to user
+                tokenDisplay.innerHTML = `<div class="alert alert-danger">
+                    <strong>Error:</strong> ${error.message || 'Error updating results. Please try again.'}
+                </div>`;
+                
+                // Reset stats on error
+                tokenCount.textContent = '0';
+                scoredTokens.textContent = '0';
+                finalScore.textContent = '0.00';
+                pValue.textContent = '0.500';
+            }
+        }
+
+        // Increase debounce timeout and ensure it's properly cleared
+        textarea.addEventListener('input', function() {
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500); // Increased to 500ms
+        });
+
+        // Add input event listeners for parameter fields to trigger updates
+        seedInput.addEventListener('input', function() {
+            const value = this.value === '' ? '' : parseInt(this.value);
+            if (isNaN(value) && this.value !== '') {
+                this.value = "0";
+            }
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500);
+        });
+
+        ngramInput.addEventListener('input', function() {
+            const value = this.value === '' ? '' : parseInt(this.value);
+            if (isNaN(value) && this.value !== '') {
+                this.value = "1";
+            }
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500);
+        });
+
+        deltaInput.addEventListener('input', function() {
+            const value = this.value === '' ? '' : parseFloat(this.value);
+            if (isNaN(value) && this.value !== '') {
+                this.value = "2.0";
+            }
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500);
+        });
+
+        temperatureInput.addEventListener('input', function() {
+            const value = this.value === '' ? '' : parseFloat(this.value);
+            if (isNaN(value) && this.value !== '') {
+                this.value = "0.8";
+            }
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500);
+        });
+
+        // Add keyboard shortcut for applying changes
+        document.addEventListener('keydown', function(e) {
+            if ((e.metaKey || e.ctrlKey) && e.key === 'Enter') {
+                e.preventDefault();
+                if (document.activeElement === promptArea) {
+                    if (generateBtn.classList.contains('generating')) {
+                        stopGeneration();
+                    } else {
+                        startGeneration();
+                    }
+                } else {
+                    applyParamsBtn.click();
+                }
+            }
+        });
+
+        detectorTypeSelect.addEventListener('change', function() {
+            if (debounceTimeout) {
+                clearTimeout(debounceTimeout);
+            }
+            debounceTimeout = setTimeout(updateTokenization, 500);
+        });
+
+        // Ensure the modal apply button properly triggers an update
+        applyParamsBtn.addEventListener('click', function() {
+            updateTokenization().then(() => {
+                const modal = bootstrap.Modal.getInstance(document.getElementById('paramsModal'));
+                if (modal) {
+                    modal.hide();
+                }
+            }).catch(error => {
+                console.error('Error applying parameters:', error);
+            });
+        });
+
+        // Initial tokenization with error handling
+        document.addEventListener('DOMContentLoaded', function() {
+            updateTokenization().catch(error => {
+                console.error('Error during initial tokenization:', error);
+            });
+        });
+
+        // Add this helper function for formatting p-values
+        function formatPValue(value) {
+            if (value >= 0.001) {
+                return value.toFixed(3);
+            } else {
+                return value.toExponential(2);
+            }
+        }
+    </script>
+</body>
+</html>

wm_interactive/web/app.py

@@ -0,0 +1,241 @@
+"""
+Main Flask application for the watermark detection web interface.
+"""
+
+from flask import Flask, render_template, request, jsonify, Response, stream_with_context
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import json
+
+from ..core.detector import MarylandDetector, MarylandDetectorZ, OpenaiDetector, OpenaiDetectorZ
+from ..core.generator import WmGenerator, OpenaiGenerator, MarylandGenerator
+from .utils import get_token_details, template_prompt
+
+CACHE_DIR = "wm_interactive/static/hf_cache"
+
+def convert_nan_to_null(obj):
+    """Convert NaN values to null for JSON serialization"""
+    import math
+    if isinstance(obj, float) and math.isnan(obj):
+        return None
+    elif isinstance(obj, dict):
+        return {k: convert_nan_to_null(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_nan_to_null(item) for item in obj]
+    return obj
+
+def set_to_int(value, default_value = None):
+    try:
+        return int(value)
+    except (ValueError, TypeError):
+        return default_value
+
+def create_detector(detector_type, tokenizer, **kwargs):
+    """Create a detector instance based on the specified type."""
+    detector_map = {
+        'maryland': MarylandDetector,
+        'marylandz': MarylandDetectorZ,
+        'openai': OpenaiDetector,
+        'openaiz': OpenaiDetectorZ
+    }
+    
+    # Validate and set default values for parameters
+    if 'seed' in kwargs:
+        kwargs['seed'] = set_to_int(kwargs['seed'], default_value = 0)
+    if 'ngram' in kwargs:
+        kwargs['ngram'] = set_to_int(kwargs['ngram'], default_value = 1)
+            
+    detector_class = detector_map.get(detector_type, MarylandDetector)
+    return detector_class(tokenizer=tokenizer, **kwargs)
+
+def create_app():
+    app = Flask(__name__, 
+                static_folder='../static',
+                template_folder='../templates')
+
+    # Add zip to Jinja's global context
+    app.jinja_env.globals.update(zip=zip)
+
+    # Pick a model
+    # model_id = "meta-llama/Llama-3.2-1B-Instruct"
+    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
+    tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir=CACHE_DIR)
+    model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir=CACHE_DIR).to("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Create default generator
+    generator = MarylandGenerator(model, tokenizer, ngram=1, seed=0)
+
+    @app.route("/", methods=["GET"])
+    def index():
+        return render_template("index.html")
+
+    @app.route("/tokenize", methods=["POST"])
+    def tokenize():
+        try:
+            data = request.get_json()
+            if not data:
+                return jsonify({'error': 'No JSON data received'}), 400
+                
+            text = data.get('text', '')
+            params = data.get('params', {})
+            
+            # Create a detector instance with the provided parameters
+            detector = create_detector(
+                detector_type=params.get('detector_type', 'maryland'),
+                tokenizer=tokenizer,
+                seed=params.get('seed', 0),
+                ngram=params.get('ngram', 1)
+            )
+            
+            if text:
+                try:
+                    display_info = get_token_details(text, detector)
+
+                    # Extract summary stats (last item in display_info)
+                    stats = display_info.pop()
+
+                    response_data = {
+                        'token_count': len(display_info),
+                        'tokens': [info['token'] for info in display_info],
+                        'colors': [info['color'] for info in display_info],
+                        'scores': [info['score'] if info.get('is_scored', False) else None for info in display_info],
+                        'pvalues': [info['pvalue'] if info.get('is_scored', False) else None for info in display_info],
+                        'final_score': stats.get('final_score', 0) if stats.get('final_score') is not None else 0,
+                        'ntoks_scored': stats.get('ntoks_scored', 0) if stats.get('ntoks_scored') is not None else 0,
+                        'final_pvalue': stats.get('final_pvalue', 0.5) if stats.get('final_pvalue') is not None else 0.5
+                    }
+                    
+                    # Convert any NaN values to null before sending
+                    response_data = convert_nan_to_null(response_data)
+                    
+                    # Ensure numeric fields have default values if they became null
+                    if response_data['final_score'] is None:
+                        response_data['final_score'] = 0
+                    if response_data['ntoks_scored'] is None:
+                        response_data['ntoks_scored'] = 0
+                    if response_data['final_pvalue'] is None:
+                        response_data['final_pvalue'] = 0.5
+                        
+                    return jsonify(response_data)
+                    
+                except Exception as e:
+                    app.logger.error(f'Error processing text: {str(e)}')
+                    return jsonify({'error': f'Error processing text: {str(e)}'}), 500
+            
+            return jsonify({
+                'token_count': 0,
+                'tokens': [],
+                'colors': [],
+                'scores': [],
+                'pvalues': [],
+                'final_score': 0,
+                'ntoks_scored': 0,
+                'final_pvalue': 0.5
+            })
+            
+        except Exception as e:
+            app.logger.error(f'Server error: {str(e)}')
+            return jsonify({'error': f'Server error: {str(e)}'}), 500
+
+    @app.route("/generate", methods=["POST"])
+    def generate():
+        try:
+            data = request.get_json()
+            if not data:
+                return jsonify({'error': 'No JSON data received'}), 400
+                
+            prompt = template_prompt(data.get('prompt', ''))
+            params = data.get('params', {})
+            temperature = float(params.get('temperature', 0.8))
+            
+            def generate_stream():
+                try:
+                    # Create generator with correct parameters
+                    generator_class = OpenaiGenerator if params.get('detector_type') == 'openai' else MarylandGenerator
+                    generator = generator_class(
+                        model=model,
+                        tokenizer=tokenizer,
+                        ngram=set_to_int(params.get('ngram', 1)),
+                        seed=set_to_int(params.get('seed', 0)),
+                        delta=float(params.get('delta', 2.0)),
+                    )
+
+                    # Get special tokens to filter out
+                    special_tokens = {
+                        '<|im_start|>', '<|im_end|>',
+                        tokenizer.pad_token, tokenizer.eos_token,
+                        tokenizer.bos_token if hasattr(tokenizer, 'bos_token') else None,
+                        tokenizer.sep_token if hasattr(tokenizer, 'sep_token') else None
+                    }
+                    special_tokens = {t for t in special_tokens if t is not None}
+
+                    # Encode prompt
+                    prompt_tokens = tokenizer.encode(prompt)
+                    prompt_size = len(prompt_tokens)
+                    max_gen_len = 100
+                    total_len = min(getattr(model.config, 'max_position_embeddings', 2048), max_gen_len + prompt_size)
+
+                    # Initialize generation
+                    tokens = torch.full((1, total_len), model.config.pad_token_id).to(model.device).long()
+                    tokens[0, :prompt_size] = torch.tensor(prompt_tokens).long()
+                    input_text_mask = tokens != model.config.pad_token_id
+
+                    # Generate token by token
+                    prev_pos = 0
+                    outputs = None  # Initialize outputs to None
+                    for cur_pos in range(prompt_size, total_len):
+                        # Get model outputs
+                        outputs = model.forward(
+                            tokens[:, prev_pos:cur_pos], 
+                            use_cache=True, 
+                            past_key_values=outputs.past_key_values if prev_pos > 0 else None
+                        )
+
+                        # Sample next token using the generator's sampling method
+                        ngram_tokens = tokens[0, cur_pos-generator.ngram:cur_pos].tolist()
+                        aux = {
+                            'ngram_tokens': ngram_tokens,
+                            'cur_pos': cur_pos,
+                        }
+                        next_token = generator.sample_next(
+                            outputs.logits[:, -1, :],
+                            aux,
+                            temperature=temperature,
+                            top_p=0.9
+                        )
+                        # Check for EOS token
+                        if next_token == model.config.eos_token_id:
+                            break
+
+                        # Decode and check if it's a special token
+                        new_text = tokenizer.decode([next_token])
+                        if new_text not in special_tokens and not any(st in new_text for st in special_tokens):
+                            yield f"data: {json.dumps({'token': new_text, 'done': False})}\n\n"
+
+                        # Update token and position
+                        tokens[0, cur_pos] = next_token
+                        prev_pos = cur_pos
+
+                    # Send final complete text, filtering out special tokens
+                    final_tokens = tokens[0, prompt_size:cur_pos+1].tolist()
+                    final_text = tokenizer.decode(final_tokens)
+                    for st in special_tokens:
+                        final_text = final_text.replace(st, '')
+                    yield f"data: {json.dumps({'text': final_text, 'done': True})}\n\n"
+                    
+                except Exception as e:
+                    app.logger.error(f'Error generating text: {str(e)}')
+                    yield f"data: {json.dumps({'error': str(e)})}\n\n"
+            
+            return Response(stream_with_context(generate_stream()), mimetype='text/event-stream')
+            
+        except Exception as e:
+            app.logger.error(f'Server error: {str(e)}')
+            return jsonify({'error': f'Server error: {str(e)}'}), 500
+
+    return app
+
+app = create_app()
+
+if __name__ == "__main__":
+    app.run(host='0.0.0.0', port=7860) 

wm_interactive/web/utils.py

@@ -0,0 +1,83 @@
+import random
+import numpy as np
+
+from ..core.detector import WmDetector
+
+def generate_pastel_color():
+    """Generate a pastel color in HSL format."""
+    h = random.random()  # Random hue
+    s = 0.3 + random.random() * 0.2  # Saturation between 0.3-0.5
+    l = 0.8 + random.random() * 0.1  # Lightness between 0.8-0.9
+    return f"hsl({h*360}, {s*100}%, {l*100}%)"
+
+def color_from_score(score: float):
+    """
+    Take a score between 0 and 1 and output the color.
+    If the score is nan, returns a pastel gray color
+    If the score is close to 0, return pastel red, if the score is close to 1 returns pastel green.
+    """
+    if isinstance(score, float) and not np.isnan(score):
+        # Red for low scores, green for high scores
+        h = 0 if score < 0.5 else 120  # 0 = red, 120 = green
+        s = 0.3 + 0.2 * abs(2 * score - 1)  # Higher saturation for extreme values
+        l = 0.85  # Keep lightness constant for pastel colors
+        return f"hsl({h}, {s*100}%, {l*100}%)"
+    return "hsl(0, 0%, 85%)"  # Pastel gray for NaN
+
+def get_token_details(
+    text: str, 
+    detector: WmDetector
+) -> tuple:
+    """
+    Run the detector on the text and outputs everything needed for display
+    """
+    # Get scores for each token
+    token_details = detector.get_details(text)
+
+    # Get p-values for each token
+    pvalues, aux_info = detector.get_pvalues_by_tok(token_details)
+    
+    display_info = []
+    for token_detail, pvalue in zip(token_details, pvalues):
+        score = token_detail['score'] if token_detail['is_scored'] else float('nan')
+        # Convert numpy types to native Python types
+        if isinstance(score, (np.floating, np.integer)):
+            score = float(score)
+        if isinstance(pvalue, (np.floating, np.integer)):
+            pvalue = float(pvalue)
+            
+        display_info.append({
+            'is_scored': token_detail['is_scored'],
+            'token': token_detail['token_text'],
+            'color': color_from_score(score),
+            'score': score,
+            'pvalue': pvalue
+        })
+    
+    # Add summary statistics and convert numpy types to native Python types
+    display_info.append({
+        'final_score': float(aux_info['final_score']),
+        'ntoks_scored': int(aux_info['ntoks_scored']),
+        'final_pvalue': float(aux_info['final_pvalue'])
+    })
+    
+    return display_info
+
+def template_prompt(instruction: str, prompt_type: str = "smollm") -> str:
+    """Template a prompt according to the model's format.
+    
+    Args:
+        instruction: The raw prompt/instruction to template
+        prompt_type: Type of prompt format (smollm, alpaca)
+
+    Returns:
+        The formatted prompt ready for the model
+    """
+    if prompt_type == "alpaca":
+        return instruction
+    elif prompt_type == "smollm":
+        prompt = "<|im_start|>system\nYou are a helpful AI assistant named SmolLM, trained by Hugging Face<|im_end|>\n"
+        prompt += f"<|im_start|>user\n{instruction}<|im_end|>\n<|im_start|>assistant\n"
+        return prompt
+    else:
+        raise ValueError(f"Prompt type {prompt_type} not supported")