fixed args

demo_watermark.py

@@ -157,7 +157,7 @@ def parse_args():
     args = parser.parse_args()
     return args
 
-def load_model():
+def load_model(args):
     args.is_seq2seq_model = any([(model_type in args.model_name_or_path) for model_type in ["t5","T0"]])
     args.is_decoder_only_model = any([(model_type in args.model_name_or_path) for model_type in ["gpt","opt","bloom"]])
     if args.is_seq2seq_model:
@@ -178,7 +178,7 @@ def load_model():
 
     return model, tokenizer, device
 
-def generate(prompt, args, model=None, tokenizer=None):
+def generate(prompt, args, model=None, device=None, tokenizer=None):
 
     print(f"Generating with {args}")
 
@@ -261,7 +261,7 @@ def detect(input_text, args, device=None, tokenizer=None):
 
 def run_gradio(args, model=None, device=None, tokenizer=None):
 
-    generate_partial = partial(generate, model=model, tokenizer=tokenizer)
+    generate_partial = partial(generate, model=model, device=None, tokenizer=tokenizer)
     detect_partial = partial(detect, device=device, tokenizer=tokenizer)
 
     with gr.Blocks() as demo:
@@ -447,9 +447,19 @@ def main(args):
         print("Prompt:")
         print(input_text)
 
-        _, _, decoded_output_without_watermark, decoded_output_with_watermark, _ = generate(input_text, args)
-        without_watermark_detection_result = detect(decoded_output_without_watermark, args)
-        with_watermark_detection_result = detect(decoded_output_with_watermark, args)
+        _, _, decoded_output_without_watermark, decoded_output_with_watermark, _ = generate(input_text, 
+                                                                                            args, 
+                                                                                            model=model, 
+                                                                                            device=device, 
+                                                                                            tokenizer=tokenizer)
+        without_watermark_detection_result = detect(decoded_output_without_watermark, 
+                                                    args, 
+                                                    device=device, 
+                                                    tokenizer=tokenizer)
+        with_watermark_detection_result = detect(decoded_output_with_watermark, 
+                                                 args, 
+                                                 device=device, 
+                                                 tokenizer=tokenizer)
 
         print("#"*term_width)
         print("Output without watermark:")