Improve UI and reduce repetitiveness of generation

api.py

@@ -19,7 +19,7 @@ async def encrypt_api(
     body: EncryptionBody,
 ):
     model, tokenizer = ModelFactory.load_model(body.gen_model)
-    text, msg_rate = generate(
+    text, msg_rate, tokens_info = generate(
         tokenizer=tokenizer,
         model=model,
         prompt=body.prompt,
@@ -32,8 +32,9 @@ async def encrypt_api(
         private_key=body.private_key,
         max_new_tokens_ratio=body.max_new_tokens_ratio,
         num_beams=body.num_beams,
+        repetition_penalty=body.repetition_penalty,
     )
-    return {"text": text, "msg_rate": msg_rate}
+    return {"text": text, "msg_rate": msg_rate, "tokens_info": tokens_info}
 
 
 @app.post("/decrypt")
@@ -78,6 +79,9 @@ async def default_config():
                     "encrypt.default", "max_new_tokens_ratio"
                 ),
                 "num_beams": GlobalConfig.get("encrypt.default", "num_beams"),
+                "repetition_penalty": GlobalConfig.get(
+                    "encrypt.default", "repetition_penalty"
+                ),
             },
             "decrypt": {
                 "gen_model": GlobalConfig.get("encrypt.default", "gen_model"),
@@ -101,9 +105,14 @@ async def default_config():
 
 
 if __name__ == "__main__":
+    # The following are mainly used to satisfy the linter
+    host = GlobalConfig.get("server", "host")
+    host = str(host) if host is not None else "0.0.0.0"
+
     port = GlobalConfig.get("server", "port")
-    if port is None:
-        port = 8000
-    else:
-        port = int(port)
-    uvicorn.run("api:app", host="0.0.0.0", port=port, workers=4)
+    port = int(port) if port is not None else 8000
+
+    workers = GlobalConfig.get("server", "workers")
+    workers = int(workers) if workers is not None else 1
+
+    uvicorn.run("api:app", host=host, port=port, workers=workers)

config.ini

@@ -1,5 +1,7 @@
 [server]
-port = int:42069
+host = str:0.0.0.0
+port = int:6969
+workers = int:4
 
 [models.names]
 gpt2 = str:openai-community/gpt2
@@ -32,6 +34,7 @@ window_length = int:1
 private_key = int:0
 max_new_tokens_ratio = float:2.0
 num_beams = int:4
+repetition_penalty = float:1.0
 
 [decrypt.default]
 gen_model = str:gpt2

demo.py

@@ -19,9 +19,10 @@ def enc_fn(
     private_key: int,
     max_new_tokens_ratio: float,
     num_beams: int,
+    repetition_penalty: float,
 ):
     model, tokenizer = ModelFactory.load_model(gen_model)
-    text, msg_rate = generate(
+    text, msg_rate, tokens_info = generate(
         tokenizer=tokenizer,
         model=model,
         prompt=prompt,
@@ -34,8 +35,32 @@ def enc_fn(
         private_key=private_key,
         max_new_tokens_ratio=max_new_tokens_ratio,
         num_beams=num_beams,
+        repetition_penalty=repetition_penalty,
     )
-    return text, msg_rate
+    highlight_base = []
+    for token in tokens_info:
+        stat = None
+        if token["base_msg"] != -1:
+            if token["base_msg"] == token["base_enc"]:
+                stat = "correct"
+            else:
+                stat = "wrong"
+        highlight_base.append((repr(token["token"])[1:-1], stat))
+
+    highlight_byte = []
+    for i, token in enumerate(tokens_info):
+        if i == 0 or tokens_info[i - 1]["byte_id"] != token["byte_id"]:
+            stat = None
+            if token["byte_msg"] != -1:
+                if token["byte_msg"] == token["byte_enc"]:
+                    stat = "correct"
+                else:
+                    stat = "wrong"
+            highlight_byte.append([repr(token["token"])[1:-1], stat])
+        else:
+            highlight_byte[-1][0] += repr(token["token"])[1:-1]
+
+    return text, highlight_base, highlight_byte, round(msg_rate * 100, 2)
 
 
 def dec_fn(
@@ -89,6 +114,7 @@ if __name__ == "__main__":
                 )
             ),
             gr.Number(int(GlobalConfig.get("encrypt.default", "num_beams"))),
+            gr.Number(float(GlobalConfig.get("encrypt.default", "repetition_penalty"))),
         ],
         outputs=[
             gr.Textbox(
@@ -96,6 +122,18 @@ if __name__ == "__main__":
                 show_label=True,
                 show_copy_button=True,
             ),
+            gr.HighlightedText(
+                label="Text containing message (Base highlighted)",
+                combine_adjacent=False,
+                show_legend=True,
+                color_map={"correct": "green", "wrong": "red"},
+            ),
+            gr.HighlightedText(
+                label="Text containing message (Byte highlighted)",
+                combine_adjacent=False,
+                show_legend=True,
+                color_map={"correct": "green", "wrong": "red"},
+            ),
             gr.Number(label="Percentage of message in text", show_label=True),
         ],
     )

main.py

@@ -171,7 +171,7 @@ def main(args):
         print(f"  Max New Tokens Ratio: {args.max_new_tokens_ratio}")
         print(f"  Number of Beams: {args.num_beams}")
         print("=" * os.get_terminal_size().columns)
-        text, msg_rate = generate(
+        text, msg_rate, tokens_info = generate(
             tokenizer=tokenizer,
             model=model,
             prompt=args.prompt,

model_factory.py

@@ -70,3 +70,10 @@ class ModelFactory:
     @classmethod
     def get_models_names(cls):
         return list(cls.models_names.keys())
+
+    @classmethod
+    def get_model_max_length(cls, name: str):
+        if name in cls.tokenizers:
+            return cls.tokenizers[name].model_max_length
+        else:
+            return 0

processors.py

@@ -127,6 +127,7 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
         self.raw_msg = msg
         self.msg = bytes_to_base(msg, self.msg_base)
         self.gamma = gamma
+        self.tokenizer = tokenizer
         special_tokens = [
             tokenizer.bos_token_id,
             tokenizer.eos_token_id,
@@ -169,20 +170,69 @@ class EncryptorLogitsProcessor(LogitsProcessor, BaseProcessor):
     def get_message_len(self):
         return len(self.msg)
 
+    def __map_input_ids(self, input_ids: torch.Tensor, base_arr, byte_arr):
+        byte_enc_msg = [-1 for _ in range(input_ids.size(0))]
+        base_enc_msg = [-1 for _ in range(input_ids.size(0))]
+        base_msg = [-1 for _ in range(input_ids.size(0))]
+        byte_msg = [-1 for _ in range(input_ids.size(0))]
+
+        values_per_byte = get_values_per_byte(self.msg_base)
+        start = self.start_pos % values_per_byte
+
+        for i, b in enumerate(base_arr):
+            base_enc_msg[i] = base_arr[i]
+            byte_enc_msg[i] = byte_arr[(i - start) // values_per_byte]
+
+        for i, b in enumerate(self.msg):
+            base_msg[i + self.start_pos] = b
+            byte_msg[i + self.start_pos] = self.raw_msg[i // values_per_byte]
+
+        return base_msg, byte_msg, base_enc_msg, byte_enc_msg
+
     def validate(self, input_ids_batch: torch.Tensor):
         res = []
+        tokens_infos = []
         for input_ids in input_ids_batch:
-            values = []
-            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)
+            # Initialization
+            base_arr = []
+
+            # Loop and obtain values of all tokens
+            for i in range(0, input_ids.size(0)):
+                base_arr.append(self._get_value(input_ids[: i + 1]))
+
+            values_per_byte = get_values_per_byte(self.msg_base)
+
+            # Transform the values to bytes
+            start = self.start_pos % values_per_byte
+            byte_arr = base_to_bytes(base_arr[start:], self.msg_base)
+
+            # Construct the
             cnt = 0
-            for i in range(len(self.raw_msg)):
+            enc_msg = byte_arr[self.start_pos // values_per_byte :]
+            for i in range(min(len(enc_msg), len(self.raw_msg))):
                 if self.raw_msg[i] == enc_msg[i]:
                     cnt += 1
             res.append(cnt / len(self.raw_msg))
 
-        return res
+            base_msg, byte_msg, base_enc_msg, byte_enc_msg = (
+                self.__map_input_ids(input_ids, base_arr, byte_arr)
+            )
+            tokens = []
+            input_strs = [self.tokenizer.decode([input]) for input in input_ids]
+            for i in range(len(base_enc_msg)):
+                tokens.append(
+                    {
+                        "token": input_strs[i],
+                        "base_enc": base_enc_msg[i],
+                        "byte_enc": byte_enc_msg[i],
+                        "base_msg": base_msg[i],
+                        "byte_msg": byte_msg[i],
+                        "byte_id": (i - start) // values_per_byte,
+                    }
+                )
+            tokens_infos.append(tokens)
+
+        return res, tokens_infos
 
 
 class DecryptorProcessor(BaseProcessor):
@@ -199,7 +249,7 @@ class DecryptorProcessor(BaseProcessor):
             bytes_msg = []
             for i, input_ids in enumerate(input_ids_batch):
                 msg.append(list())
-                for j in range(self.window_length + shift, len(input_ids)):
+                for j in range(shift, len(input_ids)):
                     # TODO: this could be slow. Considering reimplement this.
                     value = self._get_value(input_ids[: j + 1])
                     msg[i].append(value)

schemes.py

@@ -20,6 +20,7 @@ class EncryptionBody(BaseModel):
         "encrypt.default", "max_new_tokens_ratio"
     )
     num_beams: int = GlobalConfig.get("encrypt.default", "num_beams")
+    repetition_penalty: float = GlobalConfig.get('encrypt.default', "repetition_penalty")
 
 class DecryptionBody(BaseModel):
     text: str

stegno.py

@@ -20,6 +20,7 @@ def generate(
     private_key: Union[int, None] = None,
     max_new_tokens_ratio: float = 2,
     num_beams: int = 4,
+    repetition_penalty: float = 1.0,
 ):
     """
     Generate the sequence containing the hidden data.
@@ -61,17 +62,20 @@ def generate(
         salt_key=salt_key,
         private_key=private_key,
     )
-    min_length = start_pos + logits_processor.get_message_len()
-    max_length = int(
+    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
     )
+    max_length = min(max_length, tokenizer.model_max_length)
+    min_length = min(min_length, max_length)
     output_tokens = model.generate(
         **tokenized_input,
         logits_processor=transformers.LogitsProcessorList([logits_processor]),
-        min_new_tokens=min_length,
-        max_new_tokens=max_length,
+        min_length=min_length,
+        max_length=max_length,
         do_sample=True,
-        num_beams=num_beams
+        num_beams=num_beams,
+        repetition_penalty=float(repetition_penalty),
     )
 
     output_tokens = output_tokens[:, prompt_size:]
@@ -81,9 +85,9 @@ def generate(
     output_tokens_post = tokenizer(output_text, return_tensors="pt").to(
         model.device
     )
-    msg_rates = logits_processor.validate(output_tokens_post.input_ids)
+    msg_rates, tokens_infos = logits_processor.validate(output_tokens_post.input_ids)
 
-    return output_text, msg_rates[0]
+    return output_text, msg_rates[0], tokens_infos[0]
 
 
 def decrypt(