Finish baseline

main.py

@@ -0,0 +1,143 @@
+import os
+from argparse import ArgumentParser
+
+import torch
+
+from stegno import generate, decrypt
+from utils import load_model
+
+
+def create_args():
+    parser = ArgumentParser()
+
+    # Generative model
+    parser.add_argument(
+        "--gen-model",
+        type=str,
+        default="openai-community/gpt2",
+        help="Generative model (LLM) used to generate text",
+    )
+    parser.add_argument(
+        "--device", type=str, default="cpu", help="Device to load LLM"
+    )
+    # Stenography params
+    parser.add_argument(
+        "--gamma",
+        type=float,
+        default=2.0,
+        help="Bias added to scores of tokens in valid list",
+    )
+    parser.add_argument(
+        "--msg-base",
+        type=int,
+        default=2,
+        help="Base of message",
+    )
+    parser.add_argument(
+        "--seed-scheme",
+        type=str,
+        required=True,
+        help="Scheme used to compute the seed",
+    )
+    parser.add_argument(
+        "--window-length",
+        type=int,
+        default=1,
+        help="Length of window to compute the seed",
+    )
+    parser.add_argument(
+        "--salt-key", type=str, default="", help="Path to salt key"
+    )
+    parser.add_argument(
+        "--private-key", type=str, default="", help="Path to private key"
+    )
+    # Input
+    parser.add_argument(
+        "--msg", type=str, required=True, help="Path to file containing message"
+    )
+    parser.add_argument(
+        "--prompt", type=str, required=True, help="Prompt used to generate text"
+    )
+    # Mode
+    parser.add_argument(
+        "--encrypt",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--decrypt",
+        action="store_true",
+    )
+
+    return parser.parse_args()
+
+
+def main(args):
+    args.device = torch.device(args.device)
+    model, tokenizer = load_model(args.gen_model, args.device)
+
+    if os.path.isfile(args.salt_key):
+        with open(args.salt_key, "r") as f:
+            salt_key = int(f.readline())
+    else:
+        salt_key = None
+
+    if os.path.isfile(args.private_key):
+        with open(args.private_key, "r") as f:
+            private_key = int(f.readline())
+    else:
+        private_key = None
+
+    if args.encrypt:
+        if os.path.isfile(args.msg):
+            with open(args.msg, "rb") as f:
+                msg = f.read()
+        else:
+            raise ValueError(f"Message file {args.msg} is not a file")
+
+        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"  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("=" * os.get_terminal_size().columns)
+        text = generate(
+            tokenizer=tokenizer,
+            model=model,
+            prompt=args.prompt,
+            msg=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,
+        )
+        args.text = text
+
+        print(f"Text contains message:\n{text}")
+
+    if args.decrypt:
+        msgs = decrypt(
+            tokenizer=tokenizer,
+            device=args.device,
+            text=args.text,
+            msg_base=args.msg_base,
+            seed_scheme=args.seed_scheme,
+            window_length=args.window_length,
+            salt_key=args.salt_key,
+            private_key=args.private_key,
+        )
+        print("Message:")
+        for s, msg in enumerate(msgs):
+            print(f"Shift {s}: {msg}")
+
+
+if __name__ == "__main__":
+    args = create_args()
+    main(args)

processors.py

@@ -0,0 +1,165 @@
+import os
+from typing import Union
+
+import torch
+from transformers import LogitsProcessor
+
+from seed_schemes import seed_scheme_factory
+from utils import bytes_to_base, base_to_bytes, get_values_per_byte
+
+
+class BaseProcessor(object):
+    def __init__(
+        self,
+        msg_base: int,
+        vocab: list[int],
+        device: torch.device,
+        seed_scheme: str,
+        window_length: int = 1,
+        salt_key: Union[int, None] = None,
+        private_key: Union[int, None] = None,
+    ):
+        """
+        Args:
+            msg_base: base of the message.
+            vocab: vocabulary list.
+            device: device to load processor.
+            seed_scheme: scheme used to compute the seed.
+            window_length: length of window to compute the seed.
+            salt_key: salt to add to the seed.
+            private_key: private key used to compute the seed.
+        """
+        # Universal parameters
+        self.msg_base = msg_base
+        self.vocab = vocab
+        self.vocab_size = len(vocab)
+        self.device = device
+
+        # Seed parameters
+        self.seed_fn = seed_scheme_factory.get(
+            seed_scheme,
+            salt_key=salt_key,
+            private_key=private_key,
+        )
+        self.window_length = window_length
+
+        # Initialize RNG
+        self.rng = torch.Generator(device=device)
+
+        # Compute the ranges of each value in base
+        self.ranges = torch.zeros((self.msg_base + 1), dtype=torch.int64)
+        chunk_size = self.vocab_size / self.msg_base
+        r = self.vocab_size % self.msg_base
+        self.ranges[1:] = chunk_size
+        self.ranges[1 : r + 1] += 1
+        self.ranges = torch.cumsum(self.ranges, dim=0)
+
+    def _seed_rng(self, input_ids: torch.Tensor):
+        """
+        Set the seed for the rng based on the current sequences.
+
+        Args:
+            input_ids: id in the input sequence.
+        """
+        seed = self.seed_fn(input_ids[-self.window_length :])
+        self.rng.manual_seed(seed)
+
+    def _get_valid_list_ids(self, input_ids: torch.Tensor, value: int):
+        """
+        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_list = vocab_perm[self.ranges[value] : self.ranges[value + 1]]
+
+        return vocab_list
+
+    def _get_value(self, input_ids: torch.Tensor):
+        """
+        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)
+
+        cur_token = input_ids[-1]
+        cur_id = (vocab_perm == cur_token).nonzero(as_tuple=True)[0]
+        value = (cur_id < self.ranges).type(torch.int).argmax().item() - 1
+
+        return value
+
+
+class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
+    def __init__(
+        self,
+        prompt_ids: torch.Tensor,
+        msg: bytes,
+        gamma: float,
+        *args,
+        **kwargs
+    ):
+        """
+        Args:
+            msg: message to hide in the text.
+            gamma: bias add to scores of token in valid list.
+        """
+        super().__init__(*args, **kwargs)
+
+        self.start_pos = []
+        for i in range(prompt_ids.size(0)):
+            self.start_pos.append(prompt_ids[i].size(0))
+        self.msg = bytes_to_base(msg, self.msg_base)
+        self.gamma = gamma
+
+    def __call__(
+        self, input_ids_batch: torch.LongTensor, scores_batch: torch.FloatTensor
+    ):
+        # If the whole message is hidden already, then just return the raw scores.
+
+        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):
+                continue
+            scores_batch[i] = self._add_bias_to_valid_list(
+                input_ids, scores_batch[i], self.msg[msg_ptr]
+            )
+
+        return scores_batch
+
+    def _add_bias_to_valid_list(
+        self, input_ids: torch.Tensor, scores: torch.Tensor, value: int
+    ):
+        """
+        Add the bias (gamma) to the valid list tokens
+        """
+        ids = self._get_valid_list_ids(input_ids, value)
+        scores[ids] = scores[ids] + self.gamma
+        return scores
+
+    def get_message_len(self):
+        return len(self.msg)
+
+
+class DecryptorProcessor(BaseProcessor):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def decrypt(self, input_ids_batch: torch.Tensor):
+        """
+        Decrypt the text sequences.
+        """
+        shift_msg = []
+        for s 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)):
+                    # TODO: this could be slow. Considering reimplement this.
+                    value = self._get_value(input_ids[: j + 1])
+                    msg[i].append(value)
+
+                bytes_msg.append(base_to_bytes(msg[i], self.msg_base))
+            shift_msg.append(bytes_msg)
+
+        return shift_msg

seed_schemes.py

@@ -0,0 +1,39 @@
+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

@@ -0,0 +1,103 @@
+from typing import Union
+
+import torch
+import transformers
+
+from processors import EncryptorLogitsProcessor, DecryptorProcessor
+
+
+def generate(
+    tokenizer,
+    model,
+    prompt: str,
+    msg: bytes,
+    gamma: float,
+    msg_base: int,
+    seed_scheme: str,
+    window_length: int = 1,
+    salt_key: Union[int, None] = None,
+    private_key: Union[int, None] = None,
+):
+    """
+    Generate the sequence containing the hidden data.
+
+    Args:
+        tokenizer: tokenizer to use.
+        model: generative model to use.
+        prompt: input prompt.
+        msg: message to hide in the text.
+        gamma: bias add to scores of token in valid list.
+        msg_base: base of the message.
+        seed_scheme: scheme used to compute the seed.
+        window_length: length of window to compute the seed.
+        salt_key: salt to add to the seed.
+        private_key: private key used to compute the seed.
+
+    """
+    tokenized_input = tokenizer(prompt, return_tensors="pt").to(model.device)
+    logits_processor = EncryptorLogitsProcessor(
+        prompt_ids=tokenized_input.input_ids,
+        msg=msg,
+        gamma=gamma,
+        msg_base=msg_base,
+        vocab=list(tokenizer.get_vocab().values()),
+        device=model.device,
+        seed_scheme=seed_scheme,
+        window_length=window_length,
+        salt_key=salt_key,
+        private_key=private_key,
+    )
+    output_tokens = model.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,
+        do_sample=True,
+        num_beams=4,
+    )
+    output_text = tokenizer.batch_decode(
+        output_tokens, skip_special_tokens=True
+    )[0]
+
+    return output_text
+
+
+def decrypt(
+    tokenizer,
+    device: torch.device,
+    text: str,
+    msg_base: int,
+    seed_scheme: str,
+    window_length: int = 1,
+    salt_key: Union[int, None] = None,
+    private_key: Union[int, None] = None,
+):
+    """
+    Extract the hidden data from the generated sequence.
+
+    Args:
+        tokenizer: tokenizer to use.
+        text: text to decode.
+        msg_base: base of the message.
+        gamma: bias added to scores of valid list.
+        seed_scheme: scheme used to compute the seed.
+        window_length: length of window to compute the seed.
+        salt_key: salt to add to the seed.
+        private_key: private key used to compute the seed.
+    """
+    tokenized_input = tokenizer(text, return_tensors="pt").to(device)
+
+    decryptor = DecryptorProcessor(
+        msg_base=msg_base,
+        vocab=list(tokenizer.get_vocab().values()),
+        device=device,
+        seed_scheme=seed_scheme,
+        window_length=window_length,
+        salt_key=salt_key,
+        private_key=private_key,
+    )
+
+    msg = decryptor.decrypt(tokenized_input.input_ids)
+
+    return msg
+

utils.py

@@ -0,0 +1,52 @@
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+def get_values_per_byte(base: int):
+    values_per_byte = 1
+    tmp = 255 // base
+    while tmp > 0:
+        values_per_byte += 1
+        tmp = tmp // base
+    return values_per_byte
+
+
+def bytes_to_base(m: bytes, base: int) -> list[int]:
+    values_per_byte = get_values_per_byte(base)
+    values = []
+    for b in m:
+        tmp = []
+        for i in range(values_per_byte):
+            tmp.append(b % base)
+            b = b // base
+        values.extend(tmp[::-1])
+
+    return values
+
+
+def base_to_bytes(values: list[int], base: int) -> bytes:
+    values_per_byte = get_values_per_byte(base)
+
+    arr = bytearray()
+
+    i = 0
+    while i < len(values):
+        tmp = 0
+        for _ in range(values_per_byte):
+            tmp = tmp * base + values[i]
+            i += 1
+            if i >= len(values):
+                break
+        arr.append(tmp)
+
+    return bytes(arr)
+
+
+def load_model(name: str, device: torch.device):
+    model = AutoModelForCausalLM.from_pretrained(name)
+    model.to(device)
+    model.eval()
+
+    tokenizer = AutoTokenizer.from_pretrained(name)
+
+    return model, tokenizer