add option to set the watermark gamma & delta from the launcher

launcher/src/main.rs

@@ -55,6 +55,10 @@ struct Args {
     otlp_endpoint: Option<String>,
     #[clap(long, env)]
     cors_allow_origin: Vec<String>,
+    #[clap(long, env)]
+    watermark_gamma: Option<f32>,
+    #[clap(long, env)]
+    watermark_delta: Option<f32>,
 }
 
 fn main() -> ExitCode {
@@ -88,6 +92,8 @@ fn main() -> ExitCode {
         json_output,
         otlp_endpoint,
         cors_allow_origin,
+        watermark_gamma,
+        watermark_delta,
     } = args;
 
     // Signal handler
@@ -243,6 +249,8 @@ fn main() -> ExitCode {
                 huggingface_hub_cache,
                 weights_cache_override,
                 disable_custom_kernels,
+                watermark_gamma,
+                watermark_delta,
                 otlp_endpoint,
                 status_sender,
                 shutdown,
@@ -414,6 +422,8 @@ fn shard_manager(
     huggingface_hub_cache: Option<String>,
     weights_cache_override: Option<String>,
     disable_custom_kernels: bool,
+    watermark_gamma: Option<f32>,
+    watermark_delta: Option<f32>,
     otlp_endpoint: Option<String>,
     status_sender: mpsc::Sender<ShardStatus>,
     shutdown: Arc<Mutex<bool>>,
@@ -494,6 +504,16 @@ fn shard_manager(
         env.push(("DISABLE_CUSTOM_KERNELS".into(), "True".into()))
     }
 
+    // Watermark Gamma
+    if let Some(watermark_gamma) = watermark_gamma {
+        env.push(("WATERMARK_GAMMA".into(), watermark_gamma.to_string().into()))
+    }
+
+    // Watermark Delta
+    if let Some(watermark_delta) = watermark_delta {
+        env.push(("WATERMARK_DELTA".into(), watermark_delta.to_string().into()))
+    }
+
     // Start process
     tracing::info!("Starting shard {rank}");
     let mut p = match Popen::create(

server/text_generation/utils/watermark.py

@@ -13,17 +13,21 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import os
 
 import torch
 from transformers import LogitsProcessor
 
+GAMMA = os.getenv("WATERMARK_GAMMA", 0.5)
+DELTA = os.getenv("WATERMARK_DELTA", 2.0)
+
 
 class WatermarkLogitsProcessor(LogitsProcessor):
     def __init__(
             self,
             vocab_size: int,
-        gamma: float = 0.5,
-        delta: float = 2.0,
+            gamma: float = GAMMA,
+            delta: float = DELTA,
             hash_key: int = 15485863,  # just a large prime number to create a rng seed with sufficient bit width
             device: str = "cpu",
     ):