Forging prompt in the function rather than in the call

app.py

@@ -35,8 +35,8 @@ tokenizer = AutoTokenizer.from_pretrained(args.model, use_auth_token=hf_token)
 
 def embed(user, max_length, window_size, method, prompt):
     uid = USERS.index(user)
-    
     watermarker = Watermarker(tokenizer=tokenizer, model=model, window_size=window_size, payload_bits=PAYLOAD_BITS)
+    prompt = get_prompt(prompt)
     watermarked_texts = watermarker.embed(key=args.key, messages=[ uid ],
                                           max_length=max_length, method=method, prompt=prompt, window_size=window_size)
     print("watermarked_texts: ", watermarked_texts)
@@ -45,8 +45,9 @@ def embed(user, max_length, window_size, method, prompt):
     
 def detect(attacked_text, window_size, method, prompt):
     watermarker = Watermarker(tokenizer=tokenizer, model=model, window_size=window_size, payload_bits=PAYLOAD_BITS)
-
+    prompt = get_prompt(prompt)
     pvalues, messages = watermarker.detect([ attacked_text ], key=args.key, method=method, prompts=[prompt])
+
     print("messages: ", messages)
     print("p-values: ", pvalues)
     user = USERS[messages[0]]
@@ -84,8 +85,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, 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")
+    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")
 
     demo.launch()