Shows percentage of message hidden in the results; Improve the UI of the command line interface

main.py

@@ -1,4 +1,5 @@
 import os
+import json
 from argparse import ArgumentParser
 
 import torch
@@ -51,18 +52,37 @@ def create_args():
     parser.add_argument(
         "--private-key", type=str, default="", help="Path to private key"
     )
+    # Generation Params
+    parser.add_argument(
+        "--num-beams",
+        type=int,
+        default=4,
+        help="Number of beams used in beam search",
+    )
+    parser.add_argument(
+        "--max-new-tokens-ratio",
+        type=float,
+        default=2,
+        help="Ratio of max new tokens to minimum tokens required to hide message",
+    )
     # Input
     parser.add_argument(
-        "--msg", type=str, required=True, help="Path to file containing message"
+        "--msg",
+        type=str,
+        default=None,
+        help="Message or path to message to be hidden",
     )
     parser.add_argument(
-        "--prompt", type=str, default=None, help="Prompt used to generate text"
+        "--prompt",
+        type=str,
+        default=None,
+        help="Prompt or path to prompt used to generate text",
     )
     parser.add_argument(
         "--text",
         type=str,
         default=None,
-        help="Text contains the hidden message",
+        help="Text or path to text containing the hidden message",
     )
     # Mode
     parser.add_argument(
@@ -89,74 +109,109 @@ def main(args):
 
     if os.path.isfile(args.salt_key):
         with open(args.salt_key, "r") as f:
-            salt_key = int(f.readline())
+            args.salt_key = int(f.readline())
+        print(f"Read salt key from {args.salt_key}")
     else:
-        salt_key = None
+        args.salt_key = int(args.salt_key) if len(args.salt_key) > 0 else None
 
     if os.path.isfile(args.private_key):
         with open(args.private_key, "r") as f:
-            private_key = int(f.readline())
+            args.private_key = int(f.readline())
+        print(f"Read private key from {args.private_key}")
     else:
-        private_key = None
+        args.private_key = (
+            int(args.private_key) if len(args.private_key) > 0 else None
+        )
 
     if args.encrypt:
         if len(args.prompt) == 0:
             raise ValueError("Prompt cannot be empty in encrypt mode")
+        if len(args.msg) == 0:
+            raise ValueError("Message cannot be empty in encrypt mode")
+
+        if os.path.isfile(args.prompt):
+            print(f"Read prompt from {args.prompt}")
+            with open(args.prompt, "r") as f:
+                args.prompt = "".join(f.readlines())
+
         if os.path.isfile(args.msg):
+            print(f"Read message from {args.msg}")
             with open(args.msg, "rb") as f:
-                msg = f.read()
+                args.msg = f.read()
         else:
-            raise ValueError(f"Message file {args.msg} is not a file")
+            args.msg = bytes(args.msg)
 
         print("=" * os.get_terminal_size().columns)
         print("Encryption Parameters:")
         print(f"  GenModel: {args.gen_model}")
-        print(f"  Prompt: {args.prompt}")
-        print(f"  Message: {msg}")
+        print(f"  Prompt:")
+        print("- " * (os.get_terminal_size().columns // 2))
+        print(args.prompt)
+        print("- " * (os.get_terminal_size().columns // 2))
+        print(f"  Message:")
+        print("- " * (os.get_terminal_size().columns // 2))
+        print(args.msg)
+        print("- " * (os.get_terminal_size().columns // 2))
         print(f"  Gamma: {args.gamma}")
         print(f"  Message Base: {args.msg_base}")
         print(f"  Seed Scheme: {args.seed_scheme}")
         print(f"  Window Length: {args.window_length}")
-        print(f"  Salt Key: {salt_key}")
-        print(f"  Private Key: {private_key}")
+        print(f"  Salt Key: {args.salt_key}")
+        print(f"  Private Key: {args.private_key}")
+        print(f"  Max New Tokens Ratio: {args.max_new_tokens_ratio}")
+        print(f"  Number of Beams: {args.num_beams}")
         print("=" * os.get_terminal_size().columns)
-        text = generate(
+        text, msg_rate = generate(
             tokenizer=tokenizer,
             model=model,
             prompt=args.prompt,
-            msg=msg,
+            msg=args.msg,
             gamma=args.gamma,
             msg_base=args.msg_base,
             seed_scheme=args.seed_scheme,
             window_length=args.window_length,
-            salt_key=salt_key,
-            private_key=private_key,
+            salt_key=args.salt_key,
+            private_key=args.private_key,
+            max_new_tokens_ratio=args.max_new_tokens_ratio,
+            num_beams=args.num_beams,
         )
-        print(f"Text contains message:\n{text}")
-
-        if os.path.isfile(args.save_file):
+        print(f"Text contains message:")
+        print("-" * (os.get_terminal_size().columns))
+        print(text)
+        print("-" * (os.get_terminal_size().columns))
+        print(f"Successfully hide {msg_rate*100:.2f} of the message")
+        print("-" * (os.get_terminal_size().columns))
+
+        if len(args.save_file) > 0:
+            os.makedirs(os.path.dirname(args.save_file), exist_ok=True)
             with open(args.save_file, "w") as f:
                 f.write(text)
-
-        args.text = text
+            print(f"Saved result to {args.save_file}")
 
     if args.decrypt:
         if len(args.text) == 0:
             raise ValueError("Text cannot be empty in decrypt mode")
+
         if os.path.isfile(args.text):
+            print(f"Read text from {args.text}")
             with open(args.text, "r") as f:
                 lines = f.readlines()
                 args.text = "".join(lines)
+
         print("=" * os.get_terminal_size().columns)
-        print("Encryption Parameters:")
+        print("Decryption Parameters:")
         print(f"  GenModel: {args.gen_model}")
-        print(f"  Text: {args.text}")
         print(f"  Message Base: {args.msg_base}")
         print(f"  Seed Scheme: {args.seed_scheme}")
         print(f"  Window Length: {args.window_length}")
-        print(f"  Salt Key: {salt_key}")
-        print(f"  Private Key: {private_key}")
+        print(f"  Salt Key: {args.salt_key}")
+        print(f"  Private Key: {args.private_key}")
+        print(f"  Text:")
+        print("- " * (os.get_terminal_size().columns // 2))
+        print(args.text)
+        print("- " * (os.get_terminal_size().columns // 2))
         print("=" * os.get_terminal_size().columns)
+
         msgs = decrypt(
             tokenizer=tokenizer,
             device=args.device,
@@ -169,7 +224,10 @@ def main(args):
         )
         print("Message:")
         for s, msg in enumerate(msgs):
-            print(f"Shift {s}: {msg}")
+            print("-" * (os.get_terminal_size().columns))
+            print(f"Shift {s}: ")
+            print(msg[0])
+        print("-" * (os.get_terminal_size().columns))
 
 
 if __name__ == "__main__":

processors.py

@@ -107,6 +107,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         self.start_pos = []
         for i in range(prompt_ids.size(0)):
             self.start_pos.append(prompt_ids[i].size(0))
+        self.raw_msg = msg
         self.msg = bytes_to_base(msg, self.msg_base)
         self.gamma = gamma
 
@@ -139,6 +140,22 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
     def get_message_len(self):
         return len(self.msg)
 
+    def validate(self, input_ids_batch: torch.Tensor):
+        res = []
+        for input_ids in input_ids_batch:
+            values = []
+            for i in range(self.start_pos[0], input_ids.size(0)):
+                values.append(self._get_value(input_ids[: i + 1]))
+            enc_msg = base_to_bytes(values, self.msg_base)
+            cnt = 0
+            for i in range(len(self.raw_msg)):
+                if self.raw_msg[i] == enc_msg[i]:
+                    cnt += 1
+            res.append(cnt / len(self.raw_msg))
+        
+
+        return res
+
 
 class DecryptorProcessor(BaseProcessor):
     def __init__(self, *args, **kwargs):

stegno.py

@@ -17,6 +17,8 @@ def generate(
     window_length: int = 1,
     salt_key: Union[int, None] = None,
     private_key: Union[int, None] = None,
+    max_new_tokens_ratio: float = 2,
+    num_beams: int = 4,
 ):
     """
     Generate the sequence containing the hidden data.
@@ -51,15 +53,19 @@ def generate(
         **tokenized_input,
         logits_processor=transformers.LogitsProcessorList([logits_processor]),
         min_new_tokens=logits_processor.get_message_len(),
-        max_new_tokens=logits_processor.get_message_len() * 2,
+        max_new_tokens=int(
+            logits_processor.get_message_len() * max_new_tokens_ratio
+        ),
         do_sample=True,
-        num_beams=4,
+        num_beams=num_beams
     )
     output_text = tokenizer.batch_decode(
         output_tokens, skip_special_tokens=True
     )[0]
+    output_tokens_post = tokenizer(output_text, return_tensors="pt")
+    msg_rates = logits_processor.validate(output_tokens_post.input_ids)
 
-    return output_text
+    return output_text, msg_rates[0]
 
 
 def decrypt(
@@ -100,4 +106,3 @@ def decrypt(
     msg = decryptor.decrypt(tokenized_input.input_ids)
 
     return msg
-