API timeout incr to 60 sec

demo_watermark.py

@@ -210,12 +210,13 @@ def load_model(args):
 
 
 from text_generation import InferenceAPIClient
+from requests.exceptions import ReadTimeout
 def generate_with_api(prompt, args):
     hf_api_key = os.environ.get("HF_API_KEY")
     if hf_api_key is None:
         raise ValueError("HF_API_KEY environment variable not set, cannot use HF API to generate text.")
 
-    client = InferenceAPIClient(args.model_name_or_path, token=hf_api_key)
+    client = InferenceAPIClient(args.model_name_or_path, token=hf_api_key, timeout=60)
     
     assert args.n_beams == 1, "HF API models do not support beam search."
     generation_params = {
@@ -226,14 +227,22 @@ def generate_with_api(prompt, args):
         generation_params["temperature"] = args.sampling_temp
         generation_params["seed"] = args.generation_seed
 
-    generation_params["watermark"] = False
-    output = client.generate(prompt, **generation_params)
-    output_text_without_watermark = output.generated_text
-
-    generation_params["watermark"] = True
-    output = client.generate(prompt, **generation_params)
-    output_text_with_watermark = output.generated_text
-
+    timeout_msg = "[Model API timeout error. Try reducing the max_new_tokens parameter or the prompt length.]"
+    try:
+        generation_params["watermark"] = False
+        output = client.generate(prompt, **generation_params)
+        output_text_without_watermark = output.generated_text
+    except ReadTimeout as e:
+        print(e)
+        output_text_without_watermark = timeout_msg
+    try:
+        generation_params["watermark"] = True
+        output = client.generate(prompt, **generation_params)
+        output_text_with_watermark = output.generated_text
+    except ReadTimeout as e:
+        print(e)
+        output_text_with_watermark = timeout_msg
+    
     return (output_text_without_watermark, 
             output_text_with_watermark)
 
@@ -737,7 +746,7 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
         generate_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
         detect_btn.click(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
         model_selector.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
-        # When the parameters change, display the update and fire detection, since some detection params dont change the model output.
+        # When the parameters change, display the update and also fire detection, since some detection params dont change the model output.
         delta.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
         gamma.change(lambda value: str(value), inputs=[session_args], outputs=[current_parameters])
         gamma.change(fn=detect_partial, inputs=[output_without_watermark,session_args,session_tokenizer], outputs=[without_watermark_detection_result,session_args,session_tokenizer])