try fp16

app.py

@@ -22,8 +22,10 @@ arg_dict = {
     'demo_public': False, 
     # 'model_name_or_path': 'facebook/opt-125m', 
     # 'model_name_or_path': 'facebook/opt-1.3b', 
-    'model_name_or_path': 'facebook/opt-2.7b', 
-    # 'model_name_or_path': 'facebook/opt-6.7b',
+    # 'model_name_or_path': 'facebook/opt-2.7b', 
+    'model_name_or_path': 'facebook/opt-6.7b',
+    'load_fp16' : True,
+    # 'load_fp16' : False,
     'prompt_max_length': None, 
     'max_new_tokens': 200, 
     'generation_seed': 123, 

demo_watermark.py

@@ -162,6 +162,12 @@ def parse_args():
         default=True,
         help="Whether to call the torch seed function before both the unwatermarked and watermarked generate calls.",
     )
+    parser.add_argument(
+        "--load_fp16",
+        type=str2bool,
+        default=False,
+        help="Whether to run model in float16 precsion.",
+    )
     args = parser.parse_args()
     return args
 
@@ -173,13 +179,19 @@ def load_model(args):
     if args.is_seq2seq_model:
         model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name_or_path)
     elif args.is_decoder_only_model:
-        model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
+        if args.load_fp16:
+            model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path,torch_dtype=torch.float16, device_map='auto')
+        else:
+            model = AutoModelForCausalLM.from_pretrained(args.model_name_or_path)
     else:
         raise ValueError(f"Unknown model type: {args.model_name_or_path}")
 
     if args.use_gpu:
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        model = model.to(device)
+        if args.load_fp16: 
+            pass
+        else: 
+            model = model.to(device)
     else:
         device = "cpu"
     model.eval()
@@ -314,8 +326,12 @@ def run_gradio(args, model=None, device=None, tokenizer=None):
 
         # Top section, greeting and instructions
         gr.Markdown("## 💧 [A Watermark for Large Language Models](https://arxiv.org/abs/2301.10226) 🔍")
-        gr.Markdown("[jwkirchenbauer/lm-watermarking![](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/jwkirchenbauer/lm-watermarking)")
-        gr.Markdown(f"Language model: {args.model_name_or_path}")
+        with gr.Row():
+            gr.Markdown("![visitor badge](https://visitor-badge.glitch.me/badge?page_id=tomg-group-umd_lm-watermarking)")
+        with gr.Row():
+            gr.Markdown("[jwkirchenbauer/lm-watermarking![](https://badgen.net/badge/icon/GitHub?icon=github&label)](https://github.com/jwkirchenbauer/lm-watermarking)")
+        with gr.Row():
+            gr.Markdown(f"Language model: {args.model_name_or_path}")
         with gr.Accordion("Understanding the output metrics",open=False):
             gr.Markdown(
             """