Handling prompt

app.py

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