Add statistical analysis

analyse.py

@@ -0,0 +1,408 @@
+import os
+import json
+import base64
+from argparse import ArgumentParser
+
+import numpy as np
+from matplotlib import pyplot as plt
+
+import torch
+from datasets import load_dataset
+from model_factory import ModelFactory
+from stegno import generate
+
+rng = torch.Generator(device="cpu")
+rng.manual_seed(0)
+
+
+def load_msgs(msg_lens: list[int], file: str | None = None):
+    msgs = None
+    if file is not None and os.path.isfile(file):
+        with open(file, "r") as f:
+            msgs = json.load(f)
+        if "readable" not in msgs and "random" not in msgs:
+            msgs = None
+        else:
+            return msgs
+
+    msgs = {
+        "readable": [],
+        "random": [],
+    }
+
+    c4_en = load_dataset("allenai/c4", "en", split="validation", streaming=True)
+    iterator = iter(c4_en)
+
+    for length in msg_lens:
+        random_msg = torch.randint(256, (length,), generator=rng)
+        base64_msg = base64.b64encode(bytes(random_msg.tolist())).decode(
+            "ascii"
+        )
+        msgs["random"].append(base64_msg)
+
+        readable_msg = next(iterator)["text"]
+        while len(readable_msg) < length:
+            readable_msg = next(iterator)["text"]
+        msgs["readable"].append(readable_msg[:length])
+
+    return msgs
+
+
+def load_prompts(n: int, min_length: int, file: str | None = None):
+    prompts = None
+    if file is not None and os.path.isfile(file):
+        with open(file, "r") as f:
+            prompts = json.load(f)
+        return prompts
+
+    prompts = []
+
+    c4_en = load_dataset("allenai/c4", "en", split="train", streaming=True)
+    iterator = iter(c4_en)
+
+    while len(prompts) < n:
+        text = next(iterator)["text"]
+        if len(text) < min_length:
+            continue
+        prompts.append(text)
+
+    return prompts
+
+
+def create_args():
+    parser = ArgumentParser()
+
+    # messages
+    parser.add_argument(
+        "--msgs-file", type=str, default=None, help="Where messages are stored"
+    )
+    parser.add_argument(
+        "--msgs-lengths",
+        nargs=3,
+        type=int,
+        help="Range of messages' lengths. This is parsed in form: <start> <end> <step>",
+    )
+    parser.add_argument(
+        "--msgs-per-length",
+        type=int,
+        default=5,
+        help="Number of messages per length",
+    )
+    # prompts
+    parser.add_argument(
+        "--prompts-file",
+        type=str,
+        default=None,
+        help="Where prompts are stored",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=500,
+        help="Number of prompts",
+    )
+    parser.add_argument(
+        "--prompt-size",
+        type=int,
+        default=50,
+        help="Size of prompts",
+    )
+    parser.add_argument(
+        "--prompts-min-length",
+        type=int,
+        default=100,
+        help="Min length of prompts",
+    )
+    # Others
+    parser.add_argument(
+        "--overwrite",
+        action="store_true",
+        help="Whether to overwrite prompts and messages files",
+    )
+
+    # Hyperparameters
+    parser.add_argument(
+        "--gen-model",
+        type=str,
+        default="gpt2",
+        help="Model used to generate",
+    )
+    parser.add_argument(
+        "--deltas",
+        nargs=3,
+        type=float,
+        help="Range of delta. This is parsed in form: <start> <end> <step>",
+    )
+    parser.add_argument(
+        "--bases",
+        nargs=3,
+        type=int,
+        help="Range of base. This is parsed in form: <start> <end> <step>",
+    )
+    parser.add_argument(
+        "--judge-model",
+        type=str,
+        default="gpt2",
+        help="Model used to compute score perplexity of generated text",
+    )
+    # Results
+    parser.add_argument(
+        "--repeat",
+        type=int,
+        default=1,
+        help="How many times to repeat for each set of parameters, prompts and messages",
+    )
+    parser.add_argument(
+        "--results-load-file",
+        type=str,
+        default=None,
+        help="Where to load results",
+    )
+    parser.add_argument(
+        "--results-save-file",
+        type=str,
+        default=None,
+        help="Where to save results",
+    )
+    parser.add_argument(
+        "--figs-dir",
+        type=str,
+        default=None,
+        help="Where to save figures",
+    )
+
+    return parser.parse_args()
+
+
+def get_results(args, prompts, msgs):
+    model, tokenizer = ModelFactory.load_model(args.gen_model)
+    results = []
+
+    for prompt in prompts[:1]:
+        for delta in np.arange(
+            args.deltas[0], args.deltas[1] + args.deltas[2], args.deltas[2]
+        ):
+            for base in np.arange(
+                args.bases[0],
+                args.bases[1] + args.bases[2],
+                args.bases[2],
+                dtype=np.int32,
+            ):
+                for k in msgs:
+                    msg_type = k
+                    for msg in msgs[k]:
+                        msg_bytes = (
+                            msg.encode("ascii")
+                            if k == "readable"
+                            else base64.b64decode(msg)
+                        )
+                        for _ in range(args.repeat):
+                            text, msg_rate, tokens_info = generate(
+                                tokenizer=tokenizer,
+                                model=model,
+                                prompt=prompt,
+                                msg=msg_bytes,
+                                start_pos_p=[0],
+                                delta=delta,
+                                msg_base=base,
+                                seed_scheme="sha_left_hash",
+                                window_length=1,
+                                private_key=0,
+                                min_new_tokens_ratio=1,
+                                max_new_tokens_ratio=2,
+                                num_beams=4,
+                                repetition_penalty=1.5,
+                                prompt_size=args.prompt_size,
+                            )
+                            results.append(
+                                {
+                                    "msg_type": msg_type,
+                                    "delta": delta.item(),
+                                    "base": base.item(),
+                                    "perplexity": ModelFactory.compute_perplexity(
+                                        args.judge_model, text
+                                    ),
+                                    "msg_rate": msg_rate,
+                                }
+                            )
+    return results
+
+
+def process_results(results, save_dir):
+    data = {
+        "perplexities": {
+            "random": {},
+            "readable": {},
+        },
+        "msg_rates": {
+            "random": {},
+            "readable": {},
+        },
+    }
+    for r in results:
+        msg_type = r["msg_type"]
+        base = r["base"]
+        delta = r["delta"]
+        msg_rate = r["msg_rate"]
+        perplexity = r["perplexity"]
+
+        if (base, delta) not in data["msg_rates"][msg_type]:
+            data["msg_rates"][msg_type][(base, delta)] = []
+        data["msg_rates"][msg_type][(base, delta)].append(msg_rate)
+
+        if (base, delta) not in data["perplexities"][msg_type]:
+            data["perplexities"][msg_type][(base, delta)] = []
+        data["perplexities"][msg_type][(base, delta)].append(perplexity)
+
+    bases = {
+        "perplexities": {
+            "random": [],
+            "readable": [],
+        },
+        "msg_rates": {
+            "random": [],
+            "readable": [],
+        },
+    }
+    deltas = {
+        "perplexities": {
+            "random": [],
+            "readable": [],
+        },
+        "msg_rates": {
+            "random": [],
+            "readable": [],
+        },
+    }
+    values = {
+        "perplexities": {
+            "random": [],
+            "readable": [],
+        },
+        "msg_rates": {
+            "random": [],
+            "readable": [],
+        },
+    }
+    base_set = set()
+    delta_set = set()
+    for metric in data:
+        for msg_type in data[metric]:
+            for k in data[metric][msg_type]:
+                s = sum(data[metric][msg_type][k])
+                cnt = len(data[metric][msg_type][k])
+                data[metric][msg_type][k] = s / cnt
+
+                bases[metric][msg_type].append(k[0])
+                deltas[metric][msg_type].append(k[1])
+                values[metric][msg_type].append(s / cnt)
+                base_set.add(k[0])
+                delta_set.add(k[1])
+    for metric in data:
+        for msg_type in data[metric]:
+            bases[metric][msg_type] = np.array(bases[metric][msg_type], dtype=np.int32)
+            deltas[metric][msg_type] = np.array(deltas[metric][msg_type], dtype=np.int32)
+            values[metric][msg_type] = np.array(values[metric][msg_type], dtype=np.float32)
+
+    os.makedirs(save_dir, exist_ok=True)
+    for metric in data:
+        for msg_type in data[metric]:
+            fig = plt.figure(dpi=300)
+            s = lambda x: 3.0 + x * (3 if metric == "msg_rates" else 0.1)
+            plt.scatter(
+                bases[metric][msg_type],
+                deltas[metric][msg_type],
+                s(values[metric][msg_type]),
+            )
+            plt.savefig(
+                os.path.join(save_dir, f"{metric}_{msg_type}_scatter.pdf"),
+                bbox_inches="tight",
+            )
+
+    os.makedirs(os.path.join(save_dir, "delta_effect"), exist_ok=True)
+    for metric in data:
+        for msg_type in data[metric]:
+            for base_value in base_set:
+                mask = bases[metric][msg_type] == base_value
+                fig = plt.figure(dpi=300)
+                s = lambda x: x / (1.0 if metric == "msg_rates" else 10.0)
+                plt.plot(
+                    deltas[metric][msg_type][mask],
+                    values[metric][msg_type][mask],
+                )
+                plt.savefig(
+                    os.path.join(save_dir, f"delta_effect/{metric}_{msg_type}_base{base_value}.pdf"),
+                    bbox_inches="tight",
+                )
+    os.makedirs(os.path.join(save_dir, "base_effect"), exist_ok=True)
+    for metric in data:
+        for msg_type in data[metric]:
+            for delta_value in delta_set:
+                mask = deltas[metric][msg_type] == delta_value
+                fig = plt.figure(dpi=300)
+                s = lambda x: x / (1.0 if metric == "msg_rates" else 10.0)
+                plt.plot(
+                    bases[metric][msg_type][mask],
+                    values[metric][msg_type][mask],
+                )
+                plt.savefig(
+                    os.path.join(save_dir, f"base_effect/{metric}_{msg_type}_delta{delta_value}.pdf"),
+                    bbox_inches="tight",
+                )
+
+
+def main(args):
+    prompts = load_prompts(
+        args.num_prompts,
+        args.prompts_min_length,
+        args.prompts_file if not args.overwrite else None,
+    )
+
+    msgs_lens = []
+    for i in np.arange(
+        args.msgs_lengths[0],
+        args.msgs_lengths[1] + args.msgs_lengths[2],
+        args.msgs_lengths[2],
+        dtype=np.int32,
+    ):
+        for _ in range(args.msgs_per_length):
+            msgs_lens.append(i)
+
+    msgs = load_msgs(
+        msgs_lens,
+        args.msgs_file if not args.overwrite else None,
+    )
+
+    if args.msgs_file:
+        if not os.path.isfile(args.msgs_file) or args.overwrite:
+            os.makedirs(os.path.dirname(args.msgs_file), exist_ok=True)
+            with open(args.msgs_file, "w") as f:
+                json.dump(msgs, f)
+            print(f"Saved messages to {args.msgs_file}")
+    if args.prompts_file:
+        if not os.path.isfile(args.prompts_file) or args.overwrite:
+            os.makedirs(os.path.dirname(args.prompts_file), exist_ok=True)
+            with open(args.prompts_file, "w") as f:
+                json.dump(prompts, f)
+            print(f"Saved prompts to {args.prompts_file}")
+
+    if args.results_load_file:
+        with open(args.results_load_file, "r") as f:
+            results = json.load(f)
+    else:
+        results = get_results(args, prompts, msgs)
+
+    if args.results_save_file:
+        os.makedirs(os.path.dirname(args.results_save_file), exist_ok=True)
+        with open(args.results_save_file, "w") as f:
+            json.dump(results, f)
+        print(f"Saved results to {args.results_save_file}")
+
+    if args.figs_dir:
+        process_results(results, args.figs_dir)
+
+
+
+if __name__ == "__main__":
+    args = create_args()
+    main(args)

api.py

@@ -108,6 +108,9 @@ async def default_config():
                 "private_key": GlobalConfig.get(
                     "encrypt.default", "private_key"
                 ),
+                "min_new_tokens_ratio": GlobalConfig.get(
+                    "encrypt.default", "min_new_tokens_ratio"
+                ),
                 "max_new_tokens_ratio": GlobalConfig.get(
                     "encrypt.default", "max_new_tokens_ratio"
                 ),

config.ini

@@ -32,6 +32,7 @@ msg_base = int:2
 seed_scheme = str:sha_left_hash
 window_length = int:1
 private_key = int:0
+min_new_tokens_ratio = float:1.0
 max_new_tokens_ratio = float:2.0
 num_beams = int:4
 repetition_penalty = float:1.0

model_factory.py

@@ -63,7 +63,8 @@ class ModelFactory:
     @classmethod
     def load_model(cls, name):
         if name not in cls.models:
-            cls.__load_model(name)
+            if cls.__load_model(name) is None:
+                return None, None
 
         if name != cls.run_model and cls.run_model is not None:
             cls.models[cls.run_model].to(cls.load_device)
@@ -83,3 +84,43 @@ class ModelFactory:
             return cls.tokenizers[name].model_max_length
         else:
             return 0
+
+    @classmethod
+    def compute_perplexity(cls, model_name, text):
+        # This code is copied from https://huggingface.co/docs/transformers/perplexity
+        model, tokenizer = cls.load_model(model_name)
+        if model is None or tokenizer is None:
+            return 0
+        device = model.device
+        encodings = tokenizer(text, return_tensors="pt").to(device)
+
+        max_length = model.config.n_positions
+        stride = max_length//2
+        seq_len = encodings.input_ids.size(1)
+
+        nlls = []
+        prev_end_loc = 0
+        for begin_loc in range(0, seq_len, stride):
+            end_loc = min(begin_loc + max_length, seq_len)
+            trg_len = end_loc - prev_end_loc  # may be different from stride on last loop
+            input_ids = encodings.input_ids[:, begin_loc:end_loc].to(device)
+            target_ids = input_ids.clone()
+            target_ids[:, :-trg_len] = -100
+
+            with torch.no_grad():
+                outputs = model(input_ids, labels=target_ids)
+
+                # loss is calculated using CrossEntropyLoss which averages over valid labels
+                # N.B. the model only calculates loss over trg_len - 1 labels, because it internally shifts the labels
+                # to the left by 1.
+                neg_log_likelihood = outputs.loss
+
+            nlls.append(neg_log_likelihood)
+
+            prev_end_loc = end_loc
+            if end_loc == seq_len:
+                break
+
+        ppl = torch.exp(torch.stack(nlls).mean()).item()
+        return ppl
+

requirements.txt

@@ -1,6 +1,7 @@
 numpy==1.26.4
 tqdm==4.66.4
 transformers==4.41.2
+datasets==2.20.0
 PyYAML==6.0.1
 scikit-learn==1.5.0
 torch==2.3.0
@@ -8,3 +9,4 @@ cryptography==42.0.8
 fastapi
 gradio
 uvicorn
+matplotlib==3.9.1

schemes.py

@@ -49,6 +49,11 @@ class EncryptionBody(BaseModel):
         title="Private key used to compute the seed for PRF",
         ge=0,
     )
+    max_new_tokens_ratio: float = Field(
+        default=GlobalConfig.get("encrypt.default", "min_new_tokens_ratio"),
+        title="Min length of generated text compared to the minimum length required to hide the message",
+        ge=1,
+    )
     max_new_tokens_ratio: float = Field(
         default=GlobalConfig.get("encrypt.default", "max_new_tokens_ratio"),
         title="Max length of generated text compared to the minimum length required to hide the message",

stegno.py

@@ -18,9 +18,11 @@ def generate(
     window_length: int = 1,
     salt_key: Union[int, None] = None,
     private_key: Union[int, None] = None,
+    min_new_tokens_ratio: float = 1,
     max_new_tokens_ratio: float = 2,
     num_beams: int = 4,
     repetition_penalty: float = 1.0,
+    prompt_size: int = -1,
 ):
     """
     Generate the sequence containing the hidden data.
@@ -36,7 +38,6 @@ def generate(
         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.
-
     """
     if len(start_pos_p) == 1:
         start_pos = start_pos_p[0]
@@ -47,9 +48,10 @@ def generate(
     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)
+    if prompt_size == -1:
+        prompt_size = tokenized_input.input_ids.size(1)
     logits_processor = EncryptorLogitsProcessor(
-        prompt_ids=tokenized_input.input_ids,
+        prompt_ids=tokenized_input.input_ids[:prompt_size],
         msg=msg,
         start_pos=start_pos,
         delta=delta,
@@ -62,14 +64,21 @@ def generate(
         salt_key=salt_key,
         private_key=private_key,
     )
-    min_length = prompt_size + start_pos + logits_processor.get_message_len()
-    max_length = prompt_size + int(
-        start_pos + logits_processor.get_message_len() * max_new_tokens_ratio
+    min_length = (
+        prompt_size
+        + start_pos
+        + logits_processor.get_message_len() * min_new_tokens_ratio
+    )
+    max_length = (
+        prompt_size
+        + start_pos
+        + logits_processor.get_message_len() * max_new_tokens_ratio
     )
     max_length = min(max_length, tokenizer.model_max_length)
     min_length = min(min_length, max_length)
     output_tokens = model.generate(
-        **tokenized_input,
+        input_ids=tokenized_input.input_ids[:, :prompt_size],
+        attention_mask=tokenized_input.attention_mask[:, :prompt_size],
         logits_processor=transformers.LogitsProcessorList([logits_processor]),
         min_length=min_length,
         max_length=max_length,
@@ -79,10 +88,12 @@ def generate(
     )
 
     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", add_special_tokens=False).to(
-        model.device
-    )
+    output_text = tokenizer.batch_decode(
+        output_tokens, skip_special_tokens=True
+    )[0]
+    output_tokens_post = tokenizer(
+        output_text, return_tensors="pt", add_special_tokens=False
+    ).to(model.device)
     msg_rates, tokens_infos = logits_processor.validate(
         output_tokens_post.input_ids
     )

utils.py

@@ -55,3 +55,4 @@ def static_init(cls):
     if getattr(cls, "__static_init__", None):
         cls.__static_init__()
     return cls
+