committing broken state

app.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 import json
 import random
-from typing import Optional
+from typing import Optional, Tuple
 
 import gradio as gr
 import spaces
@@ -12,7 +12,7 @@ import transformers
 # the nature of the task (e.g., fatcual responses are lower entropy) or it could
 # be another
 
-_MODEL_IDENTIFIER = 'google/gemma-2b'
+_MODEL_IDENTIFIER = 'google/gemma-2b-it'
 _DETECTOR_IDENTIFIER = 'gg-hf/detector_2b_1.0_demo'
 
 _PROMPTS: tuple[str] = (
@@ -21,11 +21,10 @@ _PROMPTS: tuple[str] = (
     'prompt 3',
 )
 
-_CORRECT_ANSWERS: dict[str, bool] = {}
-
 _TORCH_DEVICE = (
     torch.device("cuda:0") if torch.cuda.is_available() else torch.device("cpu")
 )
+_ANSWERS = []
 
 _WATERMARK_CONFIG_DICT = dict(
     ngram_len=5,
@@ -70,7 +69,7 @@ _WATERMARK_CONFIG = transformers.generation.SynthIDTextWatermarkingConfig(
     **_WATERMARK_CONFIG_DICT
 )
 
-tokenizer = transformers.AutoTokenizer.from_pretrained(_MODEL_IDENTIFIER)
+tokenizer = transformers.AutoTokenizer.from_pretrained(_MODEL_IDENTIFIER, padding_side="left")
 tokenizer.pad_token_id = tokenizer.eos_token_id
 
 model = transformers.AutoModelForCausalLM.from_pretrained(_MODEL_IDENTIFIER)
@@ -89,6 +88,7 @@ detector_module.to(_TORCH_DEVICE)
 detector = transformers.generation.watermarking.SynthIDTextWatermarkDetector(
     detector_module=detector_module,
     logits_processor=logits_processor,
+    tokenizer=tokenizer,
 )
 
 
@@ -98,8 +98,9 @@ def generate_outputs(
   watermarking_config: Optional[
       transformers.generation.SynthIDTextWatermarkingConfig
   ] = None,
-) -> Sequence[str]:
-  tokenized_prompts = tokenizer(prompts, return_tensors='pt').to(_TORCH_DEVICE)
+) -> Tuple[Sequence[str], torch.Tensor]:
+  tokenized_prompts = tokenizer(prompts, return_tensors='pt', padding="longest").to(_TORCH_DEVICE)
+  input_length = tokenized_prompts.input_ids.shape[1]
   output_sequences = model.generate(
       **tokenized_prompts,
       watermarking_config=watermarking_config,
@@ -107,9 +108,10 @@ def generate_outputs(
       max_length=500,
       top_k=40,
   )
+  output_sequences = output_sequences[:, input_length:]
   detections = detector(output_sequences)
   print(detections)
-  return tokenizer.batch_decode(output_sequences)
+  return (tokenizer.batch_decode(output_sequences, skip_special_tokens=True), detections)
 
 
 with gr.Blocks() as demo:
@@ -236,25 +238,33 @@ with gr.Blocks() as demo:
     detect_btn = gr.Button('Detect', visible=False)
 
   def generate(*prompts):
-    standard = generate_outputs(prompts=prompts)
-    watermarked = generate_outputs(
+    standard, standard_detector = generate_outputs(prompts=prompts)
+    watermarked, watermarked_detector = generate_outputs(
         prompts=prompts,
         watermarking_config=_WATERMARK_CONFIG,
     )
-    responses = standard + watermarked
-    random.shuffle(responses)
+    upper_threshold = 0.9501
+    lower_threshold = 0.1209
+
+    def decision(score: float) -> str:
+      if score > upper_threshold:
+        return 'Watermarked'
+      elif lower_threshold < score < upper_threshold:
+        return 'Indeterminate'
+      else:
+        return 'Not watermarked'
 
-    _CORRECT_ANSWERS.update({
-        response: response in watermarked
-        for response in responses
-    })
+    responses = [(text, decision(score)) for text, score in zip(standard, standard_detector[0])]
+    responses += [(text, decision(score)) for text, score in zip(watermarked, watermarked_detector[0])]
+    random.shuffle(responses)
+    _ANSWERS = responses
 
     # Load model
     return {
         generate_btn: gr.Button(visible=False),
         generations_col: gr.Column(visible=True),
         generations_grp: gr.CheckboxGroup(
-            responses,
+            [response[0] for response in responses],
         ),
         reveal_btn: gr.Button(visible=True),
     }
@@ -269,17 +279,17 @@ with gr.Blocks() as demo:
     choices: list[str] = []
     value: list[str] = []
 
-    for response, is_watermarked in _CORRECT_ANSWERS.items():
-      if is_watermarked and response in user_selections:
-        choice = f'Correct! {response}'
-      elif not is_watermarked and response not in user_selections:
-        choice = f'Correct! {response}'
+    for (response, decision) in _ANSWERS:
+      if decision == "Watermarked":
+        value.append(choice)
+        if response in user_selections:
+          choice = f'Correct! {response}
+      elif decision == 'Indeterminate':
+        choice = f'Uncertain! {response}'
       else:
         choice = f'Incorrect. {response}'
 
       choices.append(choice)
-      if is_watermarked:
-        value.append(choice)
 
     return {
         reveal_btn: gr.Button(visible=False),