Removing prompt from the generated text

app.py

@@ -41,12 +41,12 @@ def embed(user, max_length, window_size, method, prompt):
                                           max_length=max_length, method=method, prompt=prompt)
     print("watermarked_texts: ", watermarked_texts)
 
-    return watermarked_texts[0]
+    return watermarked_texts[0].replace(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])
+    pvalues, messages = watermarker.detect([ prompt + attacked_text ], key=args.key, method=method, prompts=[prompt])
 
     print("messages: ", messages)
     print("p-values: ", pvalues)