Add restrictions and rename parameters to meet the one written in report

api.py

@@ -25,7 +25,7 @@ async def encrypt_api(
         prompt=body.prompt,
         msg=str.encode(body.msg),
         start_pos_p=[body.start_pos],
-        gamma=body.gamma,
+        delta=body.delta,
         msg_base=body.msg_base,
         seed_scheme=body.seed_scheme,
         window_length=body.window_length,
@@ -64,7 +64,7 @@ async def default_config():
             "encrypt": {
                 "gen_model": GlobalConfig.get("encrypt.default", "gen_model"),
                 "start_pos": GlobalConfig.get("encrypt.default", "start_pos"),
-                "gamma": GlobalConfig.get("encrypt.default", "gamma"),
+                "delta": GlobalConfig.get("encrypt.default", "delta"),
                 "msg_base": GlobalConfig.get("encrypt.default", "msg_base"),
                 "seed_scheme": GlobalConfig.get(
                     "encrypt.default", "seed_scheme"

config.ini

@@ -22,12 +22,12 @@ opt_13b = str:facebook/opt-13b
 [models.params]
 dtype = str:bfloat16
 load_device = str:cpu
-run_device = str:cuda
+run_device = str:cpu
 
 [encrypt.default]
 gen_model = str:gpt2
 start_pos = int:0
-gamma = float:10.0
+delta = float:10.0
 msg_base = int:2
 seed_scheme = str:sha_left_hash
 window_length = int:1

demo.py

@@ -12,7 +12,7 @@ def enc_fn(
     prompt: str,
     msg: str,
     start_pos: int,
-    gamma: float,
+    delta: float,
     msg_base: int,
     seed_scheme: str,
     window_length: int,
@@ -28,7 +28,7 @@ def enc_fn(
         prompt=prompt,
         msg=str.encode(msg),
         start_pos_p=[start_pos],
-        gamma=gamma,
+        delta=delta,
         msg_base=msg_base,
         seed_scheme=seed_scheme,
         window_length=window_length,
@@ -98,7 +98,7 @@ if __name__ == "__main__":
             gr.Textbox(),
             gr.Textbox(),
             gr.Number(int(GlobalConfig.get("encrypt.default", "start_pos"))),
-            gr.Number(float(GlobalConfig.get("encrypt.default", "gamma"))),
+            gr.Number(float(GlobalConfig.get("encrypt.default", "delta"))),
             gr.Number(int(GlobalConfig.get("encrypt.default", "msg_base"))),
             gr.Dropdown(
                 value=GlobalConfig.get("encrypt.default", "seed_scheme"),

main.py

@@ -25,9 +25,9 @@ def create_args():
     )
     # Stenography params
     parser.add_argument(
-        "--gamma",
+        "--delta",
         type=float,
-        default=GlobalConfig.get("encrypt.default", "gamma"),
+        default=GlobalConfig.get("encrypt.default", "delta"),
         help="Bias added to scores of tokens in valid list",
     )
     parser.add_argument(
@@ -162,7 +162,7 @@ def main(args):
         print("- " * (os.get_terminal_size().columns // 2))
         print(args.msg)
         print("- " * (os.get_terminal_size().columns // 2))
-        print(f"  Gamma: {args.gamma}")
+        print(f"  delta: {args.delta}")
         print(f"  Message Base: {args.msg_base}")
         print(f"  Seed Scheme: {args.seed_scheme}")
         print(f"  Window Length: {args.window_length}")
@@ -177,7 +177,7 @@ def main(args):
             prompt=args.prompt,
             msg=args.msg,
             start_pos_p=args.start_pos,
-            gamma=args.gamma,
+            delta=args.delta,
             msg_base=args.msg_base,
             seed_scheme=args.seed_scheme,
             window_length=args.window_length,

processors.py

@@ -104,7 +104,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         self,
         prompt_ids: torch.Tensor,
         msg: bytes,
-        gamma: float,
+        delta: float,
         tokenizer,
         start_pos: int = 0,
         *args,
@@ -113,7 +113,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         """
         Args:
             msg: message to hide in the text.
-            gamma: bias add to scores of token in valid list.
+            delta: bias add to scores of token in valid list.
         """
         super().__init__(*args, **kwargs)
         if prompt_ids.size(0) != 1:
@@ -126,7 +126,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
 
         self.raw_msg = msg
         self.msg = bytes_to_base(msg, self.msg_base)
-        self.gamma = gamma
+        self.delta = delta
         self.tokenizer = tokenizer
         special_tokens = [
             tokenizer.bos_token_id,
@@ -158,13 +158,13 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         self, input_ids: torch.Tensor, scores: torch.Tensor, value: int
     ):
         """
-        Add the bias (gamma) to the valid list tokens
+        Add the bias (delta) to the valid list tokens
         """
         ids = torch.cat(
             [self._get_valid_list_ids(input_ids, value), self.special_tokens]
         )
 
-        scores[ids] = scores[ids] + self.gamma
+        scores[ids] = scores[ids] + self.delta
         return scores
 
     def get_message_len(self):

schemes.py

@@ -1,34 +1,88 @@
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from global_config import GlobalConfig
+from model_factory import ModelFactory
+from seed_scheme_factory import SeedSchemeFactory
+from typing import Literal
 
 
 class EncryptionBody(BaseModel):
-    prompt: str
-    msg: str
-    gen_model: str = GlobalConfig.get("encrypt.default", "gen_model")
-    start_pos: int = GlobalConfig.get("encrypt.default", "start_pos")
+    prompt: str = Field(title="Prompt used to generate text")
+    msg: str = Field(title="Message wanted to hide")
+    gen_model: Literal[tuple(ModelFactory.get_models_names())] = Field(
+        default=GlobalConfig.get("encrypt.default", "gen_model"),
+        title="LLM used to generate text",
+    )
+    start_pos: int = Field(
+        default=GlobalConfig.get("encrypt.default", "start_pos"),
+        title="Start position to encrypt the message",
+        ge=0,
+    )
 
-    gamma: float = GlobalConfig.get("encrypt.default", "gamma")
-    msg_base: int = GlobalConfig.get("encrypt.default", "msg_base")
+    delta: float = Field(
+        default=GlobalConfig.get("encrypt.default", "delta"),
+        title="Hardness parameters",
+        gt=0,
+    )
+    msg_base: int = Field(
+        default=GlobalConfig.get("encrypt.default", "msg_base"),
+        title="Base of message used in base-encoding",
+        ge=2,
+    )
 
-    seed_scheme: str = GlobalConfig.get("encrypt.default", "seed_scheme")
-    window_length: int = GlobalConfig.get(
-        "encrypt.default", "window_length"
+    seed_scheme: Literal[tuple(SeedSchemeFactory.get_schemes_name())] = Field(
+        default=GlobalConfig.get("encrypt.default", "seed_scheme"),
+        title="Scheme used to compute seed for PRF",
     )
-    private_key: int = GlobalConfig.get("encrypt.default", "private_key")
-    max_new_tokens_ratio: float = GlobalConfig.get(
-        "encrypt.default", "max_new_tokens_ratio"
+    window_length: int = Field(
+        default=GlobalConfig.get("encrypt.default", "window_length"),
+        title="Window length (context size) used to compute the seed for PRF",
+        ge=1,
+    )
+    private_key: int = Field(
+        default=GlobalConfig.get("encrypt.default", "private_key"),
+        title="Private key used to compute the seed for PRF",
+        ge=0,
+    )
+    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",
+        ge=1,
+    )
+    num_beams: int = Field(
+        default=GlobalConfig.get("encrypt.default", "num_beams"),
+        title="Number of beams used in beam search",
+        ge=1,
+    )
+
+    repetition_penalty: float = Field(
+        default=GlobalConfig.get("encrypt.default", "repetition_penalty"),
+        title="Penalty used to avoid repetition when sampling tokens",
+        ge=1,
     )
-    num_beams: int = GlobalConfig.get("encrypt.default", "num_beams")
-    repetition_penalty: float = GlobalConfig.get('encrypt.default', "repetition_penalty")
 
-class DecryptionBody(BaseModel):
-    text: str
-    gen_model: str = GlobalConfig.get("decrypt.default", "gen_model")
-    msg_base: int = GlobalConfig.get("decrypt.default", "msg_base")
 
-    seed_scheme: str = GlobalConfig.get("decrypt.default", "seed_scheme")
-    window_length: int = GlobalConfig.get(
-        "decrypt.default", "window_length"
+class DecryptionBody(BaseModel):
+    text: str = Field(title="Text containing the message")
+    gen_model: Literal[tuple(ModelFactory.get_models_names())] = Field(
+        default=GlobalConfig.get("decrypt.default", "gen_model"),
+        title="LLM used to generate text",
+    )
+    msg_base: int = Field(
+        default=GlobalConfig.get("decrypt.default", "msg_base"),
+        title="Base of message used in base-encoding",
+        ge=2,
+    )
+    seed_scheme: Literal[tuple(SeedSchemeFactory.get_schemes_name())] = Field(
+        default=GlobalConfig.get("decrypt.default", "seed_scheme"),
+        title="Scheme used to compute seed for PRF",
+    )
+    window_length: int = Field(
+        default=GlobalConfig.get("decrypt.default", "window_length"),
+        title="Window length (context size) used to compute the seed for PRF",
+        ge=1,
+    )
+    private_key: int = Field(
+        default=GlobalConfig.get("decrypt.default", "private_key"),
+        title="Private key used to compute the seed for PRF",
+        ge=0,
     )
-    private_key: int = GlobalConfig.get("decrypt.default", "private_key")

stegno.py

@@ -12,7 +12,7 @@ def generate(
     prompt: str,
     msg: bytes,
     start_pos_p: list[int],
-    gamma: float,
+    delta: float,
     msg_base: int,
     seed_scheme: str,
     window_length: int = 1,
@@ -30,7 +30,7 @@ def generate(
         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.
+        delta: 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.
@@ -52,7 +52,7 @@ def generate(
         prompt_ids=tokenized_input.input_ids,
         msg=msg,
         start_pos=start_pos,
-        gamma=gamma,
+        delta=delta,
         msg_base=msg_base,
         vocab=list(tokenizer.get_vocab().values()),
         tokenizer=tokenizer,
@@ -107,7 +107,7 @@ def decrypt(
         tokenizer: tokenizer to use.
         text: text to decode.
         msg_base: base of the message.
-        gamma: bias added to scores of valid list.
+        delta: 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.