Loading the model only once and adding LLama2 prompt

app.py

@@ -21,12 +21,16 @@ EMBED_METHODS = [ 'aaronson', 'kirchenbauer', 'sampling', 'greedy' ]
 DETECT_METHODS = [ 'aaronson', 'aaronson_simplified', 'aaronson_neyman_pearson', 'kirchenbauer']
 PAYLOAD_BITS = 2
 
+watermarker = Watermarker(modelname=args.model, window_size=window_size, payload_bits=PAYLOAD_BITS)
+DEFAULT_SYSTEM_PROMPT = """\
+You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe.  Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.
+If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information.\
+"""
+
 def embed(user, max_length, window_size, method, prompt):
     uid = USERS.index(user)
     
-    watermarker = Watermarker(modelname=args.model,
-                              window_size=window_size, payload_bits=PAYLOAD_BITS)
-
+    
     watermarked_texts = watermarker.embed(key=args.key, messages=[ uid ],
                                           max_length=max_length, method=method, prompt=prompt)
     print("watermarked_texts: ", watermarked_texts)
@@ -46,7 +50,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)
 
 with gr.Blocks() as demo:
     gr.Markdown("""# LLM generation watermarking
@@ -71,8 +79,8 @@ with gr.Blocks() as demo:
     with gr.Row():
         detection_label = gr.Label(label="Detection result")
 
-    btn1.click(fn=embed, inputs=[user, text_length, window_size, embed_method, prompt], outputs=[watermarked_text], api_name="watermark")
-    btn2.click(fn=detect, inputs=[watermarked_text, window_size, detect_method, prompt], outputs=[detection_label], api_name="detect")
+    btn1.click(fn=embed, inputs=[user, text_length, window_size, embed_method, get_prompt(prompt)], outputs=[watermarked_text], api_name="watermark")
+    btn2.click(fn=detect, inputs=[watermarked_text, window_size, detect_method, get_prompt(prompt)], outputs=[detection_label], api_name="detect")
 
     demo.launch()