init

README.md

@@ -1,13 +1,13 @@
----
-title: Unigram Watermark
-emoji: 😻
-colorFrom: purple
-colorTo: green
-sdk: gradio
-sdk_version: 4.7.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Unigram-Watermark
+emoji: 👀
+colorFrom: purple
+colorTo: indigo
+sdk: gradio
+sdk_version: 4.7.1
+app_file: app.py
+pinned: false
+license: mit
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,305 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 argparse
+from argparse import Namespace
+from pprint import pprint
+from functools import partial
+
+import numpy # for gradio hot reload
+import gradio as gr
+import pathlib
+import torch
+
+from transformers import (AutoTokenizer,
+                          AutoModelForSeq2SeqLM,
+                          AutoModelForCausalLM,
+                          LogitsProcessorList,
+                          LlamaTokenizer)
+
+from gptwm import GPTWatermarkDetector, GPTWatermarkLogitsWarper
+
+def str2bool(v):
+    """Util function for user friendly boolean flag args"""
+    if isinstance(v, bool):
+        return v
+    if v.lower() in ('yes', 'true', 't', 'y', '1'):
+        return True
+    elif v.lower() in ('no', 'false', 'f', 'n', '0'):
+        return False
+    else:
+        raise argparse.ArgumentTypeError('Boolean value expected.')
+
+def parse_args():
+    """Command line argument specification"""
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--run_gradio",type=str2bool,default=True,help="Whether to launch as a gradio demo. Set to False if not installed and want to just run the stdout version.")
+    parser.add_argument("--model_name", type=str, default="facebook/opt-125m")
+    parser.add_argument("--fraction", type=float, default=0.5)
+    parser.add_argument("--strength", type=float, default=2.0)
+    parser.add_argument("--wm_key", type=int, default=0)
+    parser.add_argument("--max_new_tokens", type=int, default=300)
+    parser.add_argument("--beam_size", type=int, default=None)
+    parser.add_argument("--top_k", type=int, default=None)
+    parser.add_argument("--top_p", type=float, default=0.9)
+    parser.add_argument("--test_min_tokens", type=int, default=200)
+    parser.add_argument("--threshold", type=float, default=6.0)
+    args = parser.parse_args()
+    return args
+
+def load_model(args):
+    """Load and return the model and tokenizer"""
+    hf_token = os.getenv('HF_TOKEN')
+    if 'llama' in args.model_name:
+        tokenizer = LlamaTokenizer.from_pretrained(args.model_name, use_auth_token=hf_token, torch_dtype=torch.float16)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name, use_auth_token=hf_token, torch_dtype=torch.float16)
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, use_auth_token=hf_token, device_map='auto')
+    model.eval()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    return model, tokenizer, device
+
+def generate(prompt, args, model=None, device=None, tokenizer=None):
+    print(f"Generating with {args}")
+
+    watermark_processor = LogitsProcessorList([GPTWatermarkLogitsWarper(fraction=args.fraction,
+                                                                        strength=args.strength,
+                                                                        vocab_size=model.config.vocab_size,
+                                                                        watermark_key=args.wm_key)])
+
+    
+    batch = tokenizer(prompt, truncation=True, return_tensors="pt").to(device)
+    num_tokens = len(batch['input_ids'][0])
+    with torch.inference_mode():
+            generate_args = {
+                **batch,
+                'output_scores': True,
+                'return_dict_in_generate': True,
+                'max_new_tokens': args.max_new_tokens,
+            }
+
+            if args.beam_size is not None:
+                generate_args['num_beams'] = args.beam_size
+            else:
+                generate_args['do_sample'] = True
+                generate_args['top_k'] = args.top_k
+                generate_args['top_p'] = args.top_p
+            
+            generate_without_watermark = partial(
+                model.generate,
+                **generate_args
+            )
+            output_without_watermark = generate_without_watermark()
+            decoded_output_without_watermark = tokenizer.batch_decode(output_without_watermark['sequences'][:, num_tokens:], skip_special_tokens=True)[0]
+            generate_with_watermark = partial(
+                model.generate,
+                logits_processor=watermark_processor, 
+                **generate_args
+            )
+            output_with_watermark = generate_with_watermark()
+            decoded_gen_text_with_wm = tokenizer.batch_decode(output_with_watermark['sequences'][:, num_tokens:], skip_special_tokens=True)[0]
+
+    return (prompt,
+            decoded_output_without_watermark, 
+            decoded_gen_text_with_wm,
+            args) 
+
+
+
+def detect_demo(input_text, args, device=None, tokenizer=None):
+
+    vocab_size = 50272 if "opt" in args.model_name else tokenizer.vocab_size
+
+    watermark_detector = GPTWatermarkDetector(fraction=args.fraction,
+                                        strength=args.strength,
+                                        vocab_size=vocab_size,
+                                        watermark_key=args.wm_key)
+    output = []
+    html_output = ["Input text is too short to test."]
+    tokens = tokenizer(input_text, add_special_tokens=False)
+    gen_tokens = tokens["input_ids"]
+    if len(gen_tokens)>= args.test_min_tokens:
+        z_score,green_tokens_mask,green_tokens,total_tokens = watermark_detector.detect(gen_tokens)
+        output.append(['z-score', f"{z_score:.3g}"])
+        output.append(['green_tokens', f"{int(green_tokens):d}"])
+        output.append(['total_tokens', f"{int(total_tokens):d}"])
+        tokenarray =[tokens.token_to_chars(i) for i in range(0,len(gen_tokens))]
+        tags = [(f'<span class="green">{input_text[word.start:word.end]}</span>' if b else f'<span class="red">{input_text[word.start:word.end]}</span>') for word, b in zip(tokenarray, green_tokens_mask)]
+        html_output = f'<p>{" ".join(tags)}</p>'
+    else:
+        print(f"Input text is too short to test.")
+    return output,html_output, args
+
+def run_gradio(args, model=None, device=None, tokenizer=None):
+    """Define and launch the gradio demo interface"""
+    css = """
+    .green {
+    color: #008000 !important;  
+    border: none;  
+    font-weight: bold; 
+    }
+    .red {
+    color: #ffad99 !important;  
+    border: none;  
+    font-weight: bold;  
+    }
+    """
+
+    generate_partial = partial(generate, model=model, device=device, tokenizer=tokenizer)
+    detect_partial = partial(detect_demo, device=device, tokenizer=tokenizer)
+
+    with gr.Blocks(css=css) as demo:
+        # Top section, greeting and instructions
+        with gr.Row():
+            with gr.Row():
+                with gr.Column(scale=9):
+                    gr.Markdown(
+                    """
+                    ## 🔍 Unigram-Watermark for AI-Generated Text 
+                    
+                    ## [Paper](https://arxiv.org/abs/2306.17439)  [GitHub](https://github.com/XuandongZhao/Unigram-Watermark)
+                    """
+                    )
+
+        with gr.Accordion("Abstract",open=True):
+            gr.Markdown(
+            """
+            We instantiate our language model watermarking with the **Unigram-Watermark**——a variant of the K-gram watermark. 
+            
+            We prove that our watermark method enjoys guaranteed generation quality, correctness in watermark detection, and is robust against text editing and paraphrasing.
+            """
+            )
+
+        gr.Markdown(f"Language model: {args.model_name}")
+
+        # Construct state for parameters, define updates and toggles
+        default_prompt = args.__dict__.pop("default_prompt")
+        session_args = gr.State(value=args)
+
+        with gr.Tab("Method"):
+            with gr.Accordion("Watermark process",open=True):
+                gr.Markdown(
+                """
+                1. Randomly partition the vocabulary into two distinct sets: the green list with $\gamma N$ tokens and the red list with the remaining tokens. 
+                2. In $\hat{M}$, the logits of the language model for the green list tokens are increased by $\delta$ while the logits for tokens in the red list remain unchanged.
+                """
+                )
+            with gr.Accordion("Detect process",open=True):
+                gr.Markdown(
+                """
+                1. Count the number of green tokens in the suspect text.
+
+                2. Normalize the test-statistic $z_{y}=(|y|_G-\gamma n) / \sqrt{n \gamma(1-\gamma)}$.
+
+                3. Make a calibrated decision on whether we think the suspect text is generated from $\hat{M}$ or not. 
+                """
+                )
+        with gr.Tab("Generate and Detect"):
+            
+            with gr.Row():
+                prompt = gr.Textbox(label=f"Prompt", interactive=True,lines=10,max_lines=10, value=default_prompt)
+            with gr.Row():
+                generate_btn = gr.Button("Generate")
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Tab("Output Without Watermark"):
+                        output_without_watermark = gr.Textbox(label="Text", interactive=False,lines=10,max_lines=10)
+                    with gr.Tab("Visualization"):# ￥
+                        html_without_watermark = gr.HTML(elem_id="html-without-watermark")
+                with gr.Column(scale=1):
+                    without_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"], interactive=False,row_count=7,col_count=2)
+            with gr.Row():
+                with gr.Column(scale=1):
+                    with gr.Tab("Output With Watermark"):
+                        output_with_watermark = gr.Textbox(label="Text", interactive=False,lines=10,max_lines=10)
+                    with gr.Tab("Visualization"):# 
+                        html_with_watermark = gr.HTML(elem_id="html-with-watermark")
+                with gr.Column(scale=1):
+                    with_watermark_detection_result = gr.Dataframe(headers=["Metric", "Value"],interactive=False,row_count=7,col_count=2)
+
+            redecoded_input = gr.Textbox(visible=False)
+            truncation_warning = gr.Number(visible=False)
+            def truncate_prompt(redecoded_input, truncation_warning, orig_prompt, args):
+                if truncation_warning:
+                    return redecoded_input + f"\n\n[Prompt was truncated before generation due to length...]", args
+                else: 
+                    return orig_prompt, args
+        
+        with gr.Tab("Detector Only"):
+            with gr.Row():
+                with gr.Column(scale=2):
+                    # detect inputbox
+                    with gr.Tab("Text to Analyze"):
+                        detection_input = gr.Textbox(label="Input", interactive=True,lines=14,max_lines=14)
+                    with gr.Tab("Visualization"):
+                        html_detection = gr.HTML(elem_id="html-detection")
+                with gr.Column(scale=1):
+                    detection_result = gr.Dataframe(headers=["Metric", "Value"], interactive=False,row_count=7,col_count=2)
+            with gr.Row():
+                    # detect
+                    detect_btn = gr.Button("Detect")
+        
+
+        
+        generate_btn.click(fn=generate_partial, inputs=[prompt,session_args], outputs=[redecoded_input, output_without_watermark, output_with_watermark,session_args])
+        # Show truncated version of prompt if truncation occurred
+        redecoded_input.change(fn=truncate_prompt, inputs=[redecoded_input,truncation_warning,prompt,session_args], outputs=[prompt,session_args])
+        # Call detection when the outputs (of the generate function) are updated
+        output_without_watermark.change(fn=detect_partial, inputs=[output_without_watermark,session_args], outputs=[without_watermark_detection_result,html_without_watermark,session_args])
+        output_with_watermark.change(fn=detect_partial, inputs=[output_with_watermark,session_args], outputs=[with_watermark_detection_result,html_with_watermark,session_args])
+        # Register main detection tab click
+        detect_btn.click(fn=detect_partial, inputs=[detection_input,session_args], outputs=[detection_result, html_detection,session_args])
+
+    
+    demo.launch()
+
+def main(args): 
+    """Run a command line version of the generation and detection operations
+        and optionally launch and serve the gradio demo"""
+    # Initial arg processing and log
+
+    model, tokenizer, device = load_model(args)
+
+    # Generate and detect, report to stdout
+    input_text = (
+     "One tank tumbled down an embankment into the Tenaru River, drowning its crew."
+     " At 23:00 on 14 September, the remnants of the Kuma battalion conducted another attack on the same portion of the Marine lines, but were repulsed. "
+     "A final \"weak\" attack by the Kuma unit on the evening of 15 September was also defeated. Oka's unit of about 650 men attacked the Marines at several locations on the west side of the Lunga perimeter."
+     " At about 04:00 on 14 September, two Japanese companies attacked positions held by the 3rd Battalion, 5th Marine Regiment (3/5) near the coast and were thrown back with heavy losses."
+     " Another Japanese company captured a small ridge somewhat inland but was then pinned down by Marine artillery fire throughout the day and took heavy losses before withdrawing on the evening of 14 September."
+     " The rest of Oka's unit failed to find the Marine lines and did not participate in the attack. "
+     "At 13:05 on 14 September, Kawaguchi led the survivors of his shattered brigade away from the ridge and deeper into the jungle, where they rested and tended to their wounded all the next day. "
+     "Kawaguchi's units were then ordered to withdraw west to the Matanikau River valley to join with Oka's unit, a march over difficult terrain."
+     " Kawaguchi's troops began the march on the morning of 16 September."
+     " Almost every soldier able to walk had to help carry the wounded. "
+     "As the march progressed, the exhausted and hungry soldiers, who had eaten their last rations on the morning before their withdrawal, began to discard their heavy equipment and then their rifles. "
+     "By the time most of them reached Oka's positions at Kokumbona five days later, only half still carried their weapons."
+     " The Kuma battalion's survivors, attempting to follow Kawaguchi's Center Body forces, became lost, wandered for three weeks in the jungle, and almost starved to death before finally reaching Kawaguchi's camp."
+    )
+
+    args.default_prompt = input_text
+
+    # Launch the app to generate and detect interactively (implements the hf space demo)
+    if args.run_gradio:
+        run_gradio(args, model=model, tokenizer=tokenizer, device=device)
+
+    return
+
+if __name__ == "__main__":
+
+    args = parse_args()
+    print(args)
+
+    main(args)

gptwm.py

@@ -0,0 +1,114 @@
+import hashlib
+from typing import List
+import numpy as np
+from scipy.stats import norm
+import torch
+from transformers import LogitsWarper
+
+
+class GPTWatermarkBase:
+    """
+    Base class for watermarking distributions with fixed-group green-listed tokens.
+
+    Args:
+        fraction: The fraction of the distribution to be green-listed.
+        strength: The strength of the green-listing. Higher values result in higher logit scores for green-listed tokens.
+        vocab_size: The size of the vocabulary.
+        watermark_key: The random seed for the green-listing.
+    """
+
+    def __init__(self, fraction: float = 0.5, strength: float = 2.0, vocab_size: int = 50257, watermark_key: int = 0):
+        rng = np.random.default_rng(self._hash_fn(watermark_key))
+        mask = np.array([True] * int(fraction * vocab_size) + [False] * (vocab_size - int(fraction * vocab_size)))
+        rng.shuffle(mask)
+        self.green_list_mask = torch.tensor(mask, dtype=torch.float32)
+        self.strength = strength
+        self.fraction = fraction
+
+    @staticmethod
+    def _hash_fn(x: int) -> int:
+        """solution from https://stackoverflow.com/questions/67219691/python-hash-function-that-returns-32-or-64-bits"""
+        x = np.int64(x)
+        return int.from_bytes(hashlib.sha256(x).digest()[:4], 'little')
+
+
+class GPTWatermarkLogitsWarper(GPTWatermarkBase, LogitsWarper):
+    """
+    LogitsWarper for watermarking distributions with fixed-group green-listed tokens.
+
+    Args:
+        fraction: The fraction of the distribution to be green-listed.
+        strength: The strength of the green-listing. Higher values result in higher logit scores for green-listed tokens.
+        vocab_size: The size of the vocabulary.
+        watermark_key: The random seed for the green-listing.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def __call__(self, input_ids: torch.Tensor, scores: torch.Tensor) -> torch.FloatTensor:
+        """Add the watermark to the logits and return new logits."""
+        watermark = self.strength * self.green_list_mask
+        new_logits = scores + watermark.to(scores.device)
+        return new_logits
+
+
+class GPTWatermarkDetector(GPTWatermarkBase):
+    """
+    Class for detecting watermarks in a sequence of tokens.
+
+    Args:
+        fraction: The fraction of the distribution to be green-listed.
+        strength: The strength of the green-listing. Higher values result in higher logit scores for green-listed tokens.
+        vocab_size: The size of the vocabulary.
+        watermark_key: The random seed for the green-listing.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @staticmethod
+    def _z_score(num_green: int, total: int, fraction: float) -> float:
+        """Calculate and return the z-score of the number of green tokens in a sequence."""
+        return (num_green - fraction * total) / np.sqrt(fraction * (1 - fraction) * total)
+    
+    @staticmethod
+    def _compute_tau(m: int, N: int, alpha: float) -> float:
+        """
+        Compute the threshold tau for the dynamic thresholding.
+
+        Args:
+            m: The number of unique tokens in the sequence.
+            N: Vocabulary size.
+            alpha: The false positive rate to control.
+        Returns:
+            The threshold tau.
+        """
+        factor = np.sqrt(1 - (m - 1) / (N - 1))
+        tau = factor * norm.ppf(1 - alpha)
+        return tau
+
+    def detect(self, sequence: List[int]) -> float:
+        """Detect the watermark in a sequence of tokens and return the z value."""
+        green_tokens = int(sum(self.green_list_mask[i] for i in sequence))
+        green_tokens_mask = []
+        for i in sequence:
+            if self.green_list_mask[i]:
+                green_tokens_mask.append(True)
+            else:
+                green_tokens_mask.append(False)
+        # self.green_tokens_mask = green_tokens_mask
+        
+        return self._z_score(green_tokens, len(sequence), self.fraction), green_tokens_mask,green_tokens,len(sequence)
+
+    def unidetect(self, sequence: List[int]) -> float:
+        """Detect the watermark in a sequence of tokens and return the z value. Just for unique tokens."""
+        sequence = list(set(sequence))
+        green_tokens = int(sum(self.green_list_mask[i] for i in sequence))
+        return self._z_score(green_tokens, len(sequence), self.fraction)
+    
+    def dynamic_threshold(self, sequence: List[int], alpha: float, vocab_size: int) -> (bool, float):
+        """Dynamic thresholding for watermark detection. True if the sequence is watermarked, False otherwise."""
+        z_score = self.unidetect(sequence)
+        tau = self._compute_tau(len(list(set(sequence))), vocab_size, alpha)
+        return z_score > tau, z_score

requirements.txt

@@ -0,0 +1,5 @@
+torch
+transformers
+scipy
+accelerate
+pathlib

run_detect.py

@@ -0,0 +1,58 @@
+import argparse
+import json
+from tqdm import tqdm
+import torch
+from transformers import AutoTokenizer, LlamaTokenizer
+from gptwm import GPTWatermarkDetector
+
+
+def main(args):
+    with open(args.input_file, 'r') as f:
+        data = [json.loads(x) for x in f.read().strip().split("\n")]
+    if 'llama' in args.model_name:
+        tokenizer = LlamaTokenizer.from_pretrained(args.model_name, torch_dtype=torch.float16)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name, torch_dtype=torch.float16)
+
+    vocab_size = 50272 if "opt" in args.model_name else tokenizer.vocab_size
+
+    detector = GPTWatermarkDetector(fraction=args.fraction,
+                                    strength=args.strength,
+                                    vocab_size=vocab_size,
+                                    watermark_key=args.wm_key)
+
+    z_score_list = []
+    for idx, cur_data in tqdm(enumerate(data), total=len(data)):
+        gen_tokens = tokenizer(cur_data['gen_completion'][0], add_special_tokens=False)["input_ids"]
+        if len(gen_tokens) >= args.test_min_tokens:
+            z_score_list.append(detector.detect(gen_tokens))
+        else:
+            print(f"Warning: sequence {idx} is too short to test.")
+
+    save_dict = {
+        'z_score': z_score_list,
+        'wm_pred': [1 if z > args.threshold else 0 for z in z_score_list]
+    }
+
+    print(save_dict)
+    with open(args.input_file.replace('.jsonl', '_z.jsonl'), 'w') as f:
+        json.dump(save_dict, f)
+
+    print('Finished!')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    # parser.add_argument("--model_name", type=str, default="facebook/opt-125m")
+    parser.add_argument("--model_name", type=str, default="decapoda-research/llama-7b-hf")
+    parser.add_argument("--fraction", type=float, default=0.5)
+    parser.add_argument("--strength", type=float, default=2.0)
+    parser.add_argument("--threshold", type=float, default=6.0)
+    parser.add_argument("--wm_key", type=int, default=0)
+    parser.add_argument("--input_file", type=str, default="./data/example_output.jsonl")
+    parser.add_argument("--test_min_tokens", type=int, default=200)
+
+    args = parser.parse_args()
+
+    main(args)

run_generate.py

@@ -0,0 +1,106 @@
+import argparse
+from tqdm import tqdm
+import json
+import torch
+import os
+from transformers import AutoModelForCausalLM, AutoTokenizer, LlamaTokenizer, LogitsProcessorList
+from gptwm import GPTWatermarkLogitsWarper
+
+
+def read_file(filename):
+    with open(filename, "r") as f:
+        return [json.loads(line) for line in f.read().strip().split("\n")]
+
+
+def write_file(filename, data):
+    with open(filename, "a") as f:
+        f.write("\n".join(data) + "\n")
+
+
+def main(args):
+    output_file = f"{args.output_dir}/{args.model_name.replace('/', '-')}_strength_{args.strength}_frac_{args.fraction}_len_{args.max_new_tokens}_num_{args.num_test}.jsonl"
+    if 'llama' in args.model_name:
+        tokenizer = LlamaTokenizer.from_pretrained(args.model_name, torch_dtype=torch.float16)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(args.model_name, torch_dtype=torch.float16)
+    model = AutoModelForCausalLM.from_pretrained(args.model_name, device_map='auto')
+    model.eval()
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    watermark_processor = LogitsProcessorList([GPTWatermarkLogitsWarper(fraction=args.fraction,
+                                                                        strength=args.strength,
+                                                                        vocab_size=model.config.vocab_size,
+                                                                        watermark_key=args.wm_key)])
+
+    data = read_file(args.prompt_file)
+    num_cur_outputs = len(read_file(output_file)) if os.path.exists(output_file) else 0
+
+    outputs = []
+
+    for idx, cur_data in tqdm(enumerate(data), total=min(len(data), args.num_test)):
+        if idx < num_cur_outputs or len(outputs) >= args.num_test:
+            continue
+
+        if "gold_completion" not in cur_data and 'targets' not in cur_data:
+            continue
+        elif "gold_completion" in cur_data:
+            prefix = cur_data['prefix']
+            gold_completion = cur_data['gold_completion']
+        else:
+            prefix = cur_data['prefix']
+            gold_completion = cur_data['targets'][0]
+
+        batch = tokenizer(prefix, truncation=True, return_tensors="pt").to(device)
+        num_tokens = len(batch['input_ids'][0])
+
+        with torch.inference_mode():
+            generate_args = {
+                **batch,
+                'logits_processor': watermark_processor,
+                'output_scores': True,
+                'return_dict_in_generate': True,
+                'max_new_tokens': args.max_new_tokens,
+            }
+
+            if args.beam_size is not None:
+                generate_args['num_beams'] = args.beam_size
+            else:
+                generate_args['do_sample'] = True
+                generate_args['top_k'] = args.top_k
+                generate_args['top_p'] = args.top_p
+
+            generation = model.generate(**generate_args)
+            gen_text = tokenizer.batch_decode(generation['sequences'][:, num_tokens:], skip_special_tokens=True)
+
+        outputs.append(json.dumps({
+            "prefix": prefix,
+            "gold_completion": gold_completion,
+            "gen_completion": gen_text
+        }))
+
+        if (idx + 1) % 10 == 0:
+            write_file(output_file, outputs)
+            outputs = []
+            break
+
+    write_file(output_file, outputs)
+    print("Finished!")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--model_name", type=str, default="facebookopt-125m")
+    # parser.add_argument("--model_name", type=str, default="decapoda-research/llama-7b-hf")
+    parser.add_argument("--fraction", type=float, default=0.5)
+    parser.add_argument("--strength", type=float, default=2.0)
+    parser.add_argument("--wm_key", type=int, default=0)
+    parser.add_argument("--prompt_file", type=str, default="./data/LFQA/inputs.jsonl")
+    parser.add_argument("--output_dir", type=str, default="./data/LFQA/")
+    parser.add_argument("--max_new_tokens", type=int, default=300)
+    parser.add_argument("--num_test", type=int, default=500)
+    parser.add_argument("--beam_size", type=int, default=None)
+    parser.add_argument("--top_k", type=int, default=None)
+    parser.add_argument("--top_p", type=float, default=0.9)
+
+    args = parser.parse_args()
+    main(args)