Test debug using mistral

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="meta-llama/Llama-2-7b-chat-hf", help='Language model')
+parser.add_argument('--model', '-m', type=str, default="mistralai/Mistral-7B-Instruct-v0.1", help='Language model')
 parser.add_argument('--key', '-k', type=int, default=42,
                     help='The seed of the pseudo random number generator')
 
@@ -41,16 +41,9 @@ 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)
-    print("prompt:", prompt)
     watermarked_texts = watermarker.embed(key=args.key, messages=[ uid ],
                                           max_length=max_length+LEN_DEFAULT_PROMPT, method=method, prompt=prompt)
-    print("===")
-    print(watermarked_text)
-    print("===")
-    print(watermarked_texts[0].split("[/INST]")[1][0])
-    print("===")
-    print(watermarked_texts[0].split("[/INST]"))
-    return watermarked_texts[0].split("[/INST]")[1][0]
+    return watermarked_texts[0].split("[/INST]")[1]
     
 def detect(attacked_text, window_size, method, prompt):
     watermarker = Watermarker(tokenizer=tokenizer, model=model, window_size=window_size, payload_bits=PAYLOAD_BITS)
@@ -67,10 +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)
+    # 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]"
 
 with gr.Blocks() as demo:
     gr.Markdown("""# LLM generation watermarking