Merge pull request #5 from trnKhanh/v1.1

api.py

@@ -45,14 +45,16 @@ async def encrypt_api(
         seed_scheme=body.seed_scheme,
         window_length=body.window_length,
         private_key=body.private_key,
+        min_new_tokens_ratio=body.min_new_tokens_ratio,
         max_new_tokens_ratio=body.max_new_tokens_ratio,
+        do_sample=body.do_sample,
         num_beams=body.num_beams,
         repetition_penalty=body.repetition_penalty,
     )
     return {
         "texts": texts,
         "msgs_rates": msgs_rates,
-        "tokens_info": tokens_infos,
+        "tokens_infos": tokens_infos,
     }
 
 

schemes.py

@@ -49,7 +49,7 @@ class EncryptionBody(BaseModel):
         title="Private key used to compute the seed for PRF",
         ge=0,
     )
-    max_new_tokens_ratio: float = Field(
+    min_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,
@@ -64,6 +64,10 @@ class EncryptionBody(BaseModel):
         title="Number of beams used in beam search",
         ge=1,
     )
+    do_sample: bool = Field(
+        default=GlobalConfig.get("encrypt.default", "do_sample"),
+        title="Whether to use greedy or sampling generating"
+    )
 
     repetition_penalty: float = Field(
         default=GlobalConfig.get("encrypt.default", "repetition_penalty"),

stegno.py

@@ -78,18 +78,19 @@ def generate(
         salt_key=salt_key,
         private_key=private_key,
     )
-    min_length = (
+    min_length = int(
         prompt_size
         + start_pos
         + logits_processor.get_message_len() * min_new_tokens_ratio
     )
-    max_length = (
+    max_length = int(
         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,
         logits_processor=transformers.LogitsProcessorList([logits_processor]),