Allow users to change the start_pos

main.py

@@ -84,6 +84,14 @@ def create_args():
         default=None,
         help="Text or path to text containing the hidden message",
     )
+    # Encryption params
+    parser.add_argument(
+        "--start-pos",
+        type=int,
+        nargs="+",
+        default=[0],
+        help="Start position to input the text (not including window length). If 2 integers are provided, choose the position randomly between the two values.",
+    )
     # Mode
     parser.add_argument(
         "--encrypt",
@@ -166,6 +174,7 @@ def main(args):
             model=model,
             prompt=args.prompt,
             msg=args.msg,
+            start_pos_p=args.start_pos,
             gamma=args.gamma,
             msg_base=args.msg_base,
             seed_scheme=args.seed_scheme,
@@ -179,7 +188,7 @@ def main(args):
         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(f"Successfully hide {msg_rate*100:.2f}% of the message")
         print("-" * (os.get_terminal_size().columns))
 
         if len(args.save_file) > 0:

processors.py

@@ -4,7 +4,7 @@ from typing import Union
 import torch
 from transformers import LogitsProcessor
 
-from seed_schemes import seed_scheme_factory
+from seed_scheme_factory import SeedSchemeFactory
 from utils import bytes_to_base, base_to_bytes, get_values_per_byte
 
 
@@ -36,15 +36,20 @@ class BaseProcessor(object):
         self.device = device
 
         # Seed parameters
-        self.seed_fn = seed_scheme_factory.get(
+        seed_fn = SeedSchemeFactory.get_instance(
             seed_scheme,
             salt_key=salt_key,
             private_key=private_key,
         )
+        if seed_fn is None:
+            raise ValueError(f'Seed scheme "{seed_scheme}" is invalid')
+        else:
+            self.seed_fn = seed_fn
+
         self.window_length = window_length
 
-        # Initialize RNG
-        self.rng = torch.Generator(device=device)
+        # Initialize RNG, always use cpu generator
+        self.rng = torch.Generator(device="cpu")
 
         # Compute the ranges of each value in base
         self.ranges = torch.zeros((self.msg_base + 1), dtype=torch.int64)
@@ -69,7 +74,9 @@ class BaseProcessor(object):
         Get ids of tokens in the valid list for the current sequences.
         """
         self._seed_rng(input_ids)
-        vocab_perm = torch.randperm(self.vocab_size, generator=self.rng)
+        vocab_perm = torch.randperm(
+            self.vocab_size, generator=self.rng, device="cpu"
+        ).to(self.device)
         vocab_list = vocab_perm[self.ranges[value] : self.ranges[value + 1]]
 
         return vocab_list
@@ -79,7 +86,9 @@ class BaseProcessor(object):
         Check whether the token is in the valid list.
         """
         self._seed_rng(input_ids[:-1])
-        vocab_perm = torch.randperm(self.vocab_size, generator=self.rng)
+        vocab_perm = torch.randperm(
+            self.vocab_size, generator=self.rng, device="cpu"
+        ).to(self.device)
 
         cur_token = input_ids[-1]
         cur_id = (vocab_perm == cur_token).nonzero(as_tuple=True)[0]
@@ -94,8 +103,9 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         prompt_ids: torch.Tensor,
         msg: bytes,
         gamma: float,
+        start_pos: int = 0,
         *args,
-        **kwargs
+        **kwargs,
     ):
         """
         Args:
@@ -103,10 +113,14 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
             gamma: bias add to scores of token in valid list.
         """
         super().__init__(*args, **kwargs)
+        if prompt_ids.size(0) != 1:
+            raise RuntimeError(
+                "EncryptorLogitsProcessor does not support multiple prompts input."
+            )
+
+        self.prompt_size = prompt_ids.size(1)
+        self.start_pos = start_pos
 
-        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
@@ -118,8 +132,8 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
 
         for i, input_ids in enumerate(input_ids_batch):
             cur_pos = input_ids.size(0)
-            msg_ptr = cur_pos - self.start_pos[0]
-            if msg_ptr >= len(self.msg):
+            msg_ptr = cur_pos - (self.prompt_size + self.start_pos)
+            if msg_ptr < 0 or msg_ptr >= len(self.msg):
                 continue
             scores_batch[i] = self._add_bias_to_valid_list(
                 input_ids, scores_batch[i], self.msg[msg_ptr]
@@ -144,7 +158,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         res = []
         for input_ids in input_ids_batch:
             values = []
-            for i in range(self.start_pos[0], input_ids.size(0)):
+            for i in range(self.start_pos, input_ids.size(0)):
                 values.append(self._get_value(input_ids[: i + 1]))
             enc_msg = base_to_bytes(values, self.msg_base)
             cnt = 0
@@ -152,7 +166,6 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
                 if self.raw_msg[i] == enc_msg[i]:
                     cnt += 1
             res.append(cnt / len(self.raw_msg))
-        
 
         return res
 
@@ -166,12 +179,12 @@ class DecryptorProcessor(BaseProcessor):
         Decrypt the text sequences.
         """
         shift_msg = []
-        for s in range(get_values_per_byte(self.msg_base)):
+        for shift in range(get_values_per_byte(self.msg_base)):
             msg = []
             bytes_msg = []
             for i, input_ids in enumerate(input_ids_batch):
                 msg.append(list())
-                for j in range(self.window_length + s, len(input_ids)):
+                for j in range(self.window_length + shift, len(input_ids)):
                     # TODO: this could be slow. Considering reimplement this.
                     value = self._get_value(input_ids[: j + 1])
                     msg[i].append(value)

requirements.txt

@@ -5,3 +5,4 @@ PyYAML==6.0.1
 scikit-learn==1.5.0
 torch==2.3.0
 cryptography==42.0.8
+fastapi

seed_schemes.py

@@ -1,39 +0,0 @@
-from typing import Union, Callable
-
-import torch
-
-
-class SeedSchemeFactory:
-    def __init__(self):
-        self.seed_scheme_dict = dict()
-
-    def register(self, name: str, seed_scheme: type):
-        """
-        Register the hash scheme by name. Hash scheme must be callable.
-
-        Args:
-            name: name of seed scheme.
-            func: seed function.
-        """
-        self.seed_scheme_dict[name] = seed_scheme
-
-    def get(self, name: str, **kwargs):
-        """
-        Get the hash scheme by name.
-
-        Args:
-            name: name of seed scheme.
-        """
-        return self.seed_scheme_dict[name](**kwargs)
-
-
-class DummyHash:
-    def __init__(self, *args, **kwargs):
-        pass
-
-    def __call__(self, input_ids: torch.Tensor):
-        return input_ids[-1].item()
-
-
-seed_scheme_factory = SeedSchemeFactory()
-seed_scheme_factory.register("dummy_hash", DummyHash)

stegno.py

@@ -11,6 +11,7 @@ def generate(
     model,
     prompt: str,
     msg: bytes,
+    start_pos_p: list[int],
     gamma: float,
     msg_base: int,
     seed_scheme: str,
@@ -36,10 +37,20 @@ def generate(
         private_key: private key used to compute the seed.
 
     """
+    if len(start_pos_p) == 1:
+        start_pos = start_pos_p[0]
+    else:
+        start_pos = torch.randint(
+            start_pos_p[0], start_pos_p[1] + 1, (1,)
+        ).item()
+    start_pos = int(start_pos) + window_length
+
     tokenized_input = tokenizer(prompt, return_tensors="pt").to(model.device)
+    prompt_size = tokenized_input.input_ids.size(1)
     logits_processor = EncryptorLogitsProcessor(
         prompt_ids=tokenized_input.input_ids,
         msg=msg,
+        start_pos=start_pos,
         gamma=gamma,
         msg_base=msg_base,
         vocab=list(tokenizer.get_vocab().values()),
@@ -49,20 +60,26 @@ def generate(
         salt_key=salt_key,
         private_key=private_key,
     )
+    min_length = start_pos + logits_processor.get_message_len()
+    max_length = int(
+        start_pos + logits_processor.get_message_len() * max_new_tokens_ratio
+    )
     output_tokens = model.generate(
         **tokenized_input,
         logits_processor=transformers.LogitsProcessorList([logits_processor]),
-        min_new_tokens=logits_processor.get_message_len(),
-        max_new_tokens=int(
-            logits_processor.get_message_len() * max_new_tokens_ratio
-        ),
+        min_new_tokens=min_length,
+        max_new_tokens=max_length,
         do_sample=True,
         num_beams=num_beams
     )
+
+    output_tokens = output_tokens[:, prompt_size:]
     output_text = tokenizer.batch_decode(
         output_tokens, skip_special_tokens=True
     )[0]
-    output_tokens_post = tokenizer(output_text, return_tensors="pt")
+    output_tokens_post = tokenizer(output_text, return_tensors="pt").to(
+        model.device
+    )
     msg_rates = logits_processor.validate(output_tokens_post.input_ids)
 
     return output_text, msg_rates[0]

test.py

@@ -0,0 +1,14 @@
+from seed_scheme_factory import SeedSchemeFactory
+import torch
+
+if __name__ == "__main__":
+    seed_fn = SeedSchemeFactory.get_instance("sha_left_hash", private_key=18)
+    rng = torch.Generator()
+    rng.manual_seed(1)
+    input_ids = torch.randint(0, 2**32, (8,), generator=rng)
+    print("input_ids =", input_ids)
+    if seed_fn is not None:
+        seed = seed_fn(input_ids)
+        print("      ", 2**64-1)
+        print("seed =", seed)
+