Back to llama

app.py

@@ -13,7 +13,7 @@ hf_token = os.getenv('HF_TOKEN')
 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
 
 parser = argparse.ArgumentParser(description='Generative Text Watermarking demo')
-parser.add_argument('--model', '-m', type=str, default="mistralai/Mistral-7B-Instruct-v0.1", help='Language model')
+parser.add_argument('--model', '-m', type=str, default="meta-llama/Llama-2-7b-chat-hf", help='Language model')
 parser.add_argument('--key', '-k', type=int, default=42,
                     help='The seed of the pseudo random number generator')
 
@@ -60,11 +60,11 @@ def detect(attacked_text, window_size, method, prompt):
     return label
 
 def get_prompt(message: str) -> str:
-    # texts = [f'<s>[INST] <<SYS>>\n{DEFAULT_SYSTEM_PROMPT}\n<</SYS>>\n\n']
-    # # The first user input is _not_ stripped
-    # texts.append(f'{message} [/INST]')
-    # return ''.join(texts)
-    return f"[INST] "+message+ " [/INST]"
+    texts = [f'<s>[INST] <<SYS>>\n{DEFAULT_SYSTEM_PROMPT}\n<</SYS>>\n\n']
+    # The first user input is _not_ stripped
+    texts.append(f'{message} [/INST]')
+    return ''.join(texts)
+
 
 with gr.Blocks() as demo:
     gr.Markdown("""# LLM generation watermarking

watermark.py

@@ -58,7 +58,7 @@ class WatermarkingAaronsonLogitsProcessor( WatermarkingLogitsProcessor):
 class WatermarkingKirchenbauerLogitsProcessor(WatermarkingLogitsProcessor):
     def __init__(self, *args,
                  gamma = 0.25,
-                 delta = 5.0,
+                 delta = 15.0,
                  **kwargs):
         super().__init__(*args, **kwargs)
         self.gamma = gamma