Integrate global config

api.py

@@ -1,66 +1,109 @@
+import base64
+
 import torch
 from fastapi import FastAPI
+import uvicorn
 
 from stegno import generate, decrypt
 from utils import load_model
+from seed_scheme_factory import SeedSchemeFactory
+from model_factory import ModelFactory
+from global_config import GlobalConfig
+from schemes import DecryptionBody, EncryptionBody
 
 app = FastAPI()
 
 
-@app.get("/encrypt")
-async def encrypt(
-    prompt: str,
-    msg: str,
-    gen_model: str = "openai-community/gpt2",
-    device: str = "cpu",
-    start_pos: int = 0,
-    gamma: float = 2.0,
-    msg_base: int = 2,
-    seed_scheme: str = "dummy_hash",
-    window_length: int = 1,
-    private_key: int = 0,
-    max_new_tokens_ratio: float = 2,
-    num_beams: int = 4,
+@app.post("/encrypt")
+async def encrypt_api(
+    body: EncryptionBody,
 ):
-    model, tokenizer = load_model(gen_model, torch.device(device))
+    model, tokenizer = ModelFactory.load_model(body.gen_model)
     text, msg_rate = generate(
         tokenizer=tokenizer,
         model=model,
-        prompt=prompt,
-        msg=str.encode(msg),
-        start_pos_p=[start_pos],
-        gamma=gamma,
-        msg_base=msg_base,
-        seed_scheme=seed_scheme,
-        window_length=window_length,
-        private_key=private_key,
-        max_new_tokens_ratio=max_new_tokens_ratio,
-        num_beams=num_beams,
+        prompt=body.prompt,
+        msg=str.encode(body.msg),
+        start_pos_p=[body.start_pos],
+        gamma=body.gamma,
+        msg_base=body.msg_base,
+        seed_scheme=body.seed_scheme,
+        window_length=body.window_length,
+        private_key=body.private_key,
+        max_new_tokens_ratio=body.max_new_tokens_ratio,
+        num_beams=body.num_beams,
     )
     return {"text": text, "msg_rate": msg_rate}
 
 
-@app.get("/decrypt")
-async def dec(
-    text: str,
-    gen_model: str = "openai-community/gpt2",
-    device: str = "cpu",
-    msg_base: int = 2,
-    seed_scheme: str = "dummy_hash",
-    window_length: int = 1,
-    private_key: int = 0,
-):
-    model, tokenizer = load_model(gen_model, torch.device(device))
+@app.post("/decrypt")
+async def decrypt_api(body: DecryptionBody):
+    model, tokenizer = ModelFactory.load_model(body.gen_model)
     msgs = decrypt(
         tokenizer=tokenizer,
         device=model.device,
-        text=text,
-        msg_base=msg_base,
-        seed_scheme=seed_scheme,
-        window_length=window_length,
-        private_key=private_key,
+        text=body.text,
+        msg_base=body.msg_base,
+        seed_scheme=body.seed_scheme,
+        window_length=body.window_length,
+        private_key=body.private_key,
     )
-    msg_text = ""
-    for i, msg in enumerate(msgs):
-        msg_text += f"Shift {i}: {msg}\n\n"
-    return {"msg": msg_text}
+    msg_b64 = {}
+    for i, s_msg in enumerate(msgs):
+        msg_b64[i] = []
+        for msg in s_msg:
+            msg_b64[i].append(base64.b64encode(msg))
+    return msg_b64
+
+
+@app.get("/configs")
+async def default_config():
+    configs = {
+        "default": {
+            "encrypt": {
+                "gen_model": GlobalConfig.get("encrypt.default", "gen_model"),
+                "start_pos": GlobalConfig.get("encrypt.default", "start_pos"),
+                "gamma": GlobalConfig.get("encrypt.default", "gamma"),
+                "msg_base": GlobalConfig.get("encrypt.default", "msg_base"),
+                "seed_scheme": GlobalConfig.get(
+                    "encrypt.default", "seed_scheme"
+                ),
+                "window_length": GlobalConfig.get(
+                    "encrypt.default", "window_length"
+                ),
+                "private_key": GlobalConfig.get(
+                    "encrypt.default", "private_key"
+                ),
+                "max_new_tokens_ratio": GlobalConfig.get(
+                    "encrypt.default", "max_new_tokens_ratio"
+                ),
+                "num_beams": GlobalConfig.get("encrypt.default", "num_beams"),
+            },
+            "decrypt": {
+                "gen_model": GlobalConfig.get("encrypt.default", "gen_model"),
+                "msg_base": GlobalConfig.get("encrypt.default", "msg_base"),
+                "seed_scheme": GlobalConfig.get(
+                    "encrypt.default", "seed_scheme"
+                ),
+                "window_length": GlobalConfig.get(
+                    "encrypt.default", "window_length"
+                ),
+                "private_key": GlobalConfig.get(
+                    "encrypt.default", "private_key"
+                ),
+            },
+        },
+        "seed_schemes": SeedSchemeFactory.get_schemes_name(),
+        "models": ModelFactory.get_models_names(),
+    }
+
+    return configs
+
+
+if __name__ == "__main__":
+    port = GlobalConfig.get("server", "port")
+    if port is None:
+        port = 8000
+    else:
+        port = int(port)
+    uvicorn.run(app, host="0.0.0.0", port=port)

demo.py

@@ -1,26 +1,26 @@
 import torch
 import gradio as gr
 
-from utils import load_model
+from model_factory import ModelFactory
 from stegno import generate, decrypt
 from seed_scheme_factory import SeedSchemeFactory
+from global_config import GlobalConfig
 
 
 def enc_fn(
-    gen_model: str = "openai-community/gpt2",
-    device: str = "cpu",
-    prompt: str = "",
-    msg: str = "",
-    start_pos: int = 0,
-    gamma: float = 2.0,
-    msg_base: int = 2,
-    seed_scheme: str = "dummy_hash",
-    window_length: int = 1,
-    private_key: int = 0,
-    max_new_tokens_ratio: float = 2,
-    num_beams: int = 4,
+    gen_model: str,
+    prompt: str,
+    msg: str,
+    start_pos: int,
+    gamma: float,
+    msg_base: int,
+    seed_scheme: str,
+    window_length: int,
+    private_key: int,
+    max_new_tokens_ratio: float,
+    num_beams: int,
 ):
-    model, tokenizer = load_model(gen_model, torch.device(device))
+    model, tokenizer = ModelFactory.load_model(gen_model)
     text, msg_rate = generate(
         tokenizer=tokenizer,
         model=model,
@@ -39,15 +39,14 @@ def enc_fn(
 
 
 def dec_fn(
-    gen_model: str = "openai-community/gpt2",
-    device: str = "cpu",
-    text: str = "",
-    msg_base: int = 2,
-    seed_scheme: str = "dummy_hash",
-    window_length: int = 1,
-    private_key: int = 0,
+    gen_model: str,
+    text: str,
+    msg_base: int,
+    seed_scheme: str,
+    window_length: int,
+    private_key: int,
 ):
-    model, tokenizer = load_model(gen_model, torch.device(device))
+    model, tokenizer = ModelFactory.load_model(gen_model)
     msgs = decrypt(
         tokenizer=tokenizer,
         device=model.device,
@@ -67,34 +66,56 @@ if __name__ == "__main__":
     enc = gr.Interface(
         fn=enc_fn,
         inputs=[
-            gr.Textbox("openai-community/gpt2"),
-            gr.Textbox("cpu"),
+            gr.Dropdown(
+                value=GlobalConfig.get("encrypt.default", "gen_model"),
+                choices=ModelFactory.get_models_names(),
+            ),
             gr.Textbox(),
             gr.Textbox(),
-            gr.Number(),
-            gr.Number(10.0),
-            gr.Number(2),
-            gr.Dropdown(value="dummy_hash", choices=SeedSchemeFactory.get_schemes_name()),
-            gr.Number(1),
-            gr.Number(),
-            gr.Number(2),
-            gr.Number(4),
+            gr.Number(int(GlobalConfig.get("encrypt.default", "start_pos"))),
+            gr.Number(float(GlobalConfig.get("encrypt.default", "gamma"))),
+            gr.Number(int(GlobalConfig.get("encrypt.default", "msg_base"))),
+            gr.Dropdown(
+                value=GlobalConfig.get("encrypt.default", "seed_scheme"),
+                choices=SeedSchemeFactory.get_schemes_name(),
+            ),
+            gr.Number(
+                int(GlobalConfig.get("encrypt.default", "window_length"))
+            ),
+            gr.Number(int(GlobalConfig.get("encrypt.default", "private_key"))),
+            gr.Number(
+                float(
+                    GlobalConfig.get("encrypt.default", "max_new_tokens_ratio")
+                )
+            ),
+            gr.Number(int(GlobalConfig.get("encrypt.default", "num_beams"))),
         ],
         outputs=[
-            gr.Textbox(label="Text containing message", show_label=True, show_copy_button=True),
+            gr.Textbox(
+                label="Text containing message",
+                show_label=True,
+                show_copy_button=True,
+            ),
             gr.Number(label="Percentage of message in text", show_label=True),
         ],
     )
     dec = gr.Interface(
         fn=dec_fn,
         inputs=[
-            gr.Textbox("openai-community/gpt2"),
-            gr.Textbox("cpu"),
+            gr.Dropdown(
+                value=GlobalConfig.get("decrypt.default", "gen_model"),
+                choices=ModelFactory.get_models_names(),
+            ),
             gr.Textbox(),
-            gr.Number(2),
-            gr.Dropdown(value="dummy_hash", choices=SeedSchemeFactory.get_schemes_name()),
-            gr.Number(1),
-            gr.Number(),
+            gr.Number(int(GlobalConfig.get("decrypt.default", "msg_base"))),
+            gr.Dropdown(
+                value=GlobalConfig.get("decrypt.default", "seed_scheme"),
+                choices=SeedSchemeFactory.get_schemes_name(),
+            ),
+            gr.Number(
+                int(GlobalConfig.get("decrypt.default", "window_length"))
+            ),
+            gr.Number(int(GlobalConfig.get("decrypt.default", "private_key"))),
         ],
         outputs=[
             gr.Textbox(label="Message", show_label=True),

main.py

@@ -6,6 +6,8 @@ import torch
 
 from stegno import generate, decrypt
 from utils import load_model
+from global_config import GlobalConfig
+from model_factory import ModelFactory
 
 
 def create_args():
@@ -15,7 +17,7 @@ def create_args():
     parser.add_argument(
         "--gen-model",
         type=str,
-        default="openai-community/gpt2",
+        default=GlobalConfig.get("encrypt.default", "gen_model"),
         help="Generative model (LLM) used to generate text",
     )
     parser.add_argument(
@@ -25,25 +27,25 @@ def create_args():
     parser.add_argument(
         "--gamma",
         type=float,
-        default=2.0,
+        default=GlobalConfig.get("encrypt.default", "gamma"),
         help="Bias added to scores of tokens in valid list",
     )
     parser.add_argument(
         "--msg-base",
         type=int,
-        default=2,
+        default=GlobalConfig.get("encrypt.default", "msg_base"),
         help="Base of message",
     )
     parser.add_argument(
         "--seed-scheme",
         type=str,
-        required=True,
+        default=GlobalConfig.get("encrypt.default", "seed_scheme"),
         help="Scheme used to compute the seed",
     )
     parser.add_argument(
         "--window-length",
         type=int,
-        default=1,
+        default=GlobalConfig.get("encrypt.default", "window_length"),
         help="Length of window to compute the seed",
     )
     parser.add_argument(
@@ -56,13 +58,13 @@ def create_args():
     parser.add_argument(
         "--num-beams",
         type=int,
-        default=4,
+        default=GlobalConfig.get("encrypt.default", "num_beams"),
         help="Number of beams used in beam search",
     )
     parser.add_argument(
         "--max-new-tokens-ratio",
         type=float,
-        default=2,
+        default=GlobalConfig.get("encrypt.default", "max_new_tokens_ratio"),
         help="Ratio of max new tokens to minimum tokens required to hide message",
     )
     # Input
@@ -89,7 +91,7 @@ def create_args():
         "--start-pos",
         type=int,
         nargs="+",
-        default=[0],
+        default=[GlobalConfig.get("encrypt.default", "start_pos")],
         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

processors.py

@@ -52,7 +52,9 @@ class BaseProcessor(object):
         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).to(self.device)
+        self.ranges = torch.zeros((self.msg_base + 1), dtype=torch.int64).to(
+            self.device
+        )
         chunk_size = self.vocab_size / self.msg_base
         r = self.vocab_size % self.msg_base
         self.ranges[1:] = chunk_size
@@ -103,6 +105,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         prompt_ids: torch.Tensor,
         msg: bytes,
         gamma: float,
+        tokenizer,
         start_pos: int = 0,
         *args,
         **kwargs,
@@ -124,6 +127,15 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         self.raw_msg = msg
         self.msg = bytes_to_base(msg, self.msg_base)
         self.gamma = gamma
+        special_tokens = [
+            tokenizer.bos_token_id,
+            tokenizer.eos_token_id,
+            tokenizer.sep_token_id,
+            tokenizer.pad_token_id,
+            tokenizer.cls_token_id,
+        ]
+        special_tokens = [x for x in special_tokens if x is not None]
+        self.special_tokens = torch.tensor(special_tokens, device=self.device)
 
     def __call__(
         self, input_ids_batch: torch.LongTensor, scores_batch: torch.FloatTensor
@@ -147,7 +159,10 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         """
         Add the bias (gamma) to the valid list tokens
         """
-        ids = self._get_valid_list_ids(input_ids, value)
+        ids = torch.cat(
+            [self._get_valid_list_ids(input_ids, value), self.special_tokens]
+        )
+
         scores[ids] = scores[ids] + self.gamma
         return scores
 

requirements.txt

@@ -7,3 +7,4 @@ torch==2.3.0
 cryptography==42.0.8
 fastapi
 gradio
+uvicorn

schemes.py

@@ -0,0 +1,33 @@
+from pydantic import BaseModel
+from global_config import GlobalConfig
+
+
+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")
+
+    gamma: float = GlobalConfig.get("encrypt.default", "gamma")
+    msg_base: int = GlobalConfig.get("encrypt.default", "msg_base")
+
+    seed_scheme: str = GlobalConfig.get("encrypt.default", "seed_scheme")
+    window_length: int = GlobalConfig.get(
+        "encrypt.default", "window_length"
+    )
+    private_key: int = GlobalConfig.get("encrypt.default", "private_key")
+    max_new_tokens_ratio: float = GlobalConfig.get(
+        "encrypt.default", "max_new_tokens_ratio"
+    )
+    num_beams: int = GlobalConfig.get("encrypt.default", "num_beams")
+
+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"
+    )
+    private_key: int = GlobalConfig.get("decrypt.default", "private_key")

stegno.py

@@ -54,6 +54,7 @@ def generate(
         gamma=gamma,
         msg_base=msg_base,
         vocab=list(tokenizer.get_vocab().values()),
+        tokenizer=tokenizer,
         device=model.device,
         seed_scheme=seed_scheme,
         window_length=window_length,

utils.py

@@ -50,3 +50,8 @@ def load_model(name: str, device: torch.device):
     tokenizer = AutoTokenizer.from_pretrained(name)
 
     return model, tokenizer
+
+def static_init(cls):
+    if getattr(cls, "__static_init__", None):
+        cls.__static_init__()
+    return cls