Upload 3 files

app.py

@@ -1,13 +1,18 @@
 import gradio as gr
 from models.watermark_faster import watermark_model
-import pdb
 from options import get_parser_main_model
 
 opts = get_parser_main_model().parse_args()
-model = watermark_model(language='English', mode=opts.mode, tau_word=0.8, lamda=0.83)
-def watermark_embed_demo(raw):
+model = watermark_model(language=opts.language, mode=opts.mode, tau_word=opts.tau_word, lamda=opts.lamda)
 
-    watermarked_text = model.embed(raw)
+def create_model(language,tau_word):
+    global model
+    model = watermark_model(language=language, mode=opts.mode, tau_word=tau_word, lamda=opts.lamda)
+    # gr.update(visible=True)
+    return language,tau_word
+def watermark_embed_demo(raw,tau_word):
+
+    watermarked_text = model.embed(raw,tau_word)
     return watermarked_text
 
 def watermark_extract(raw):
@@ -16,42 +21,75 @@ def watermark_extract(raw):
 
     return f"{confidence:.2f}%"
 
-def precise_watermark_detect(raw):
-    is_watermark, p_value, n, ones, z_value = model.watermark_detector_precise(raw)
+def precise_watermark_detect(raw,tau_word):
+    is_watermark, p_value, n, ones, z_value = model.watermark_detector_precise(raw,tau_word)
     confidence = (1 - p_value) * 100
 
     return f"{confidence:.2f}%"
 
-
 demo = gr.Blocks()
 with demo:
     with gr.Column():
-        gr.Markdown("# Watermarking Text Generated by Black-Box Language Models")
-
-        inputs = gr.TextArea(label="Input text", placeholder="Copy your text here...")
-        output = gr.Textbox(label="Watermarked Text")
-        analysis_button = gr.Button("Inject Watermark")
-        inputs_embed = [inputs]
-        analysis_button.click(fn=watermark_embed_demo, inputs=inputs_embed, outputs=output)
-
-        inputs_w = gr.TextArea(label="Text to Analyze", placeholder="Copy your watermarked text here...")
-
-        mode = gr.Dropdown(
-            label="Detection Mode", choices=["Fast", "Precise"], default="Fast"
-        )
-        output_detect = gr.Textbox(label="Confidence (the likelihood of the text containing a watermark)")
-        detect_button = gr.Button("Detect")
-
-        def detect_watermark(inputs_w, mode):
-            if mode == "Fast":
-                return watermark_extract(inputs_w)
-            else:
-                return precise_watermark_detect(inputs_w)
+        with gr.Row():
+            with gr.Column(scale=9):
+                gr.Markdown(
+                """
+                # 💦[Watermarking Text Generated by Black-Box Language Models](https://arxiv.org/abs/2305.08883)
+                """
+                )
+            language = gr.Dropdown(
+                label="Language", choices=["English", "Chinese"], value="English"
+            )
+            tau_word = gr.Number(label="tau_word", value=0.8)
+                
+    # with gr.Column():
+    #     with gr.Row():
+    #         gr.Markdown("# Watermarking Text Generated by Black-Box Language Models")
+    #         with gr.Row(scale=0.25):
+    #             language = gr.Dropdown(
+    #                 label="Language", choices=["English", "Chinese"], value="English"
+    #             )
+    #             tau_word = gr.Number(label="tau_word", value=0.8)#gr.Slider(0, 1, value=0.8, label="tau_word", info="Choose between 0 and 1")
+        # model_button = gr.Button("Load Model")
+        # inputs = [language,tau_word]
+        # model_button.click(fn=create_model, inputs=inputs,outputs=[language,tau_word])
+        with gr.Tab("Welcome"):
+            gr.Markdown(
+                        """
+                        This space exhibits a watermarking technique that allows third parties to independently inject an authentication watermark into generated text. 
+                        We provide implementations for both English and Chinese text (you can select the respective language in the top right corner). 
+                        Furthermore, you can adjust the value of $\\tau_{word}$ to control the similarity between the original text and the watermarked text. 
+                        We recommend setting $\\tau_{word}$ at 0.8 for English and 0.75 for Chinese. 
+                        Generally, a larger $\\tau_{word}$ increases the similarity between the original and watermarked text, but it also weakens the strength of the watermark.
+                        More details can be found in our [ArXiv preprint](https://arxiv.org/abs/2305.08883).
+                        """
+                        )
+        with gr.Tab("Watermark Injection & Detection"):
+            language.change(fn=create_model, inputs=language,outputs=language)
+            with gr.Row():
+                inputs = gr.TextArea(label="Input text", placeholder="Copy your text here...")
+                output = gr.Textbox(label="Watermarked Text",lines=7)
+            analysis_button = gr.Button("Inject Watermark")
+            inputs_embed = [inputs,tau_word]
+            analysis_button.click(fn=watermark_embed_demo, inputs=inputs_embed, outputs=output)
+            
+            inputs_w = gr.TextArea(label="Text to Analyze", placeholder="Copy your watermarked text here...")
+            with gr.Row():
+                mode = gr.Dropdown(
+                    label="Detection Mode", choices=["Fast", "Precise"], value="Fast"
+                )
+                output_detect = gr.Textbox(label="Confidence (the likelihood of the text containing a watermark)")
+            detect_button = gr.Button("Detect")
 
-        detect_button.click(fn=detect_watermark, inputs=[inputs_w, mode], outputs=output_detect)
+            def detect_watermark(inputs_w, mode, tau_word):
+                if mode == "Fast":
+                    return watermark_extract(inputs_w)
+                else:
+                    return precise_watermark_detect(inputs_w,tau_word)
 
+            detect_button.click(fn=detect_watermark, inputs=[inputs_w, mode, tau_word], outputs=output_detect)
 
 if __name__ == "__main__":
     gr.close_all()
     demo.title = "Watermarking Text Generated by Black-Box Language Models"
-    demo.launch()
+    demo.launch(share = True, server_port=8898)

models/watermark_faster.py

@@ -8,8 +8,8 @@ import hashlib
 from scipy.stats import norm
 import gensim
 import pdb
-# from transformers import BertForMaskedLM as WoBertForMaskedLM
-# from wobert import WoBertTokenizer
+from transformers import BertForMaskedLM as WoBertForMaskedLM
+from wobert import WoBertTokenizer
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 
 from transformers import BertForMaskedLM, BertTokenizer, RobertaForSequenceClassification, RobertaTokenizer
@@ -21,8 +21,6 @@ import paddle
 from jieba import posseg
 paddle.enable_static()
 import re
-nltk.download('punkt')
-nltk.download('averaged_perceptron_tagger')
 def cut_sent(para):
     para = re.sub('([。！？\?])([^”’])', r'\1\n\2', para)  
     para = re.sub('([。！？\?][”’])([^，。！？\?\n ])', r'\1\n\2', para)  
@@ -70,6 +68,7 @@ class watermark_model:
             self.w2v_model = api.load("glove-wiki-gigaword-100")
             nltk.download('stopwords')
             self.stop_words = set(stopwords.words('english'))
+            self.nlp = spacy.load('en_core_web_sm')
 
     def cut(self,ori_text,text_len):
         if self.language == 'Chinese':
@@ -273,8 +272,7 @@ class watermark_model:
         return all_processed_tokens,new_index_space
 
 
-    def filter_candidates(self, init_candidates_list, tokens, index_space, input_text):
-        
+    def filter_candidates(self, init_candidates_list, tokens, index_space, input_text, tau_word):
         all_context_word_similarity_scores = self.context_word_sim(init_candidates_list, tokens, index_space, input_text)
 
         all_sentence_similarity_scores = self.sentence_sim(init_candidates_list, tokens, index_space, input_text)
@@ -287,7 +285,7 @@ class watermark_model:
             for idx, candidate in enumerate(init_candidates):
                 global_word_similarity_score = self.global_word_sim(tokens[masked_token_index], candidate)
                 word_similarity_score = self.lamda*context_word_similarity_scores[idx]+(1-self.lamda)*global_word_similarity_score
-                if word_similarity_score >= self.tau_word and sentence_similarity_scores[idx] >= self.tau_sent:
+                if word_similarity_score >= tau_word and sentence_similarity_scores[idx] >= self.tau_sent:
                     filtered_candidates.append((candidate, word_similarity_score))
 
             if len(filtered_candidates) >= 1:
@@ -320,7 +318,7 @@ class watermark_model:
 
         return best_candidates, new_index_space
         
-    def watermark_embed(self,text):
+    def watermark_embed(self,text,tau_word):
         input_text = text
         # Tokenize the input text
         tokens = self.tokenizer.tokenize(input_text) 
@@ -344,7 +342,7 @@ class watermark_model:
         init_candidates, new_index_space = self.candidates_gen(tokens,index_space,input_text, 8, 0)
         if len(new_index_space)==0:
             return text
-        enhanced_candidates, new_index_space = self.filter_candidates(init_candidates,tokens,new_index_space,input_text)
+        enhanced_candidates, new_index_space = self.filter_candidates(init_candidates,tokens,new_index_space,input_text,tau_word)
         
         enhanced_candidates, new_index_space = self.get_candidate_encodings(tokens, enhanced_candidates, new_index_space)
         
@@ -356,7 +354,7 @@ class watermark_model:
             watermarked_text = re.sub(r'(?<=[\u4e00-\u9fff])\s+(?=[\u4e00-\u9fff，。？！、：])|(?<=[\u4e00-\u9fff，。？！、：])\s+(?=[\u4e00-\u9fff])', '', watermarked_text)
         return watermarked_text
         
-    def embed(self, ori_text):
+    def embed(self, ori_text, tau_word):
         sents = self.sent_tokenize(ori_text)
         sents = [s for s in sents if s.strip()]
         num_sents = len(sents)
@@ -369,9 +367,9 @@ class watermark_model:
                 sent_pair = sents[i]
             # keywords = jieba.analyse.extract_tags(sent_pair, topK=5, withWeight=False)
             if len(watermarked_text) == 0:
-                watermarked_text = self.watermark_embed(sent_pair)
+                watermarked_text = self.watermark_embed(sent_pair, tau_word)
             else:
-                watermarked_text = watermarked_text + self.watermark_embed(sent_pair)
+                watermarked_text = watermarked_text + self.watermark_embed(sent_pair, tau_word)
         if len(self.get_encodings_fast(ori_text)) == 0:
             # print(ori_text)
             return ''
@@ -411,7 +409,7 @@ class watermark_model:
         is_watermark = z >= threshold
         return is_watermark, p_value, n, ones, z
     
-    def get_encodings_precise(self, text):
+    def get_encodings_precise(self, text, tau_word):
         # pdb.set_trace()
         sents = self.sent_tokenize(text)
         sents = [s for s in sents if s.strip()]
@@ -441,7 +439,7 @@ class watermark_model:
                 continue
             
             init_candidates, new_index_space = self.candidates_gen(tokens,index_space,sent_pair, 8, 0)
-            enhanced_candidates, new_index_space = self.filter_candidates(init_candidates,tokens,new_index_space,sent_pair)
+            enhanced_candidates, new_index_space = self.filter_candidates(init_candidates,tokens,new_index_space,sent_pair,tau_word)
             
             # pdb.set_trace()
             for j,idx in enumerate(new_index_space):
@@ -451,9 +449,9 @@ class watermark_model:
         return encodings
 
 
-    def watermark_detector_precise(self,text,alpha=0.05):
+    def watermark_detector_precise(self,text,tau_word,alpha=0.05):
         p = 0.5
-        encodings = self.get_encodings_precise(text)
+        encodings = self.get_encodings_precise(text,tau_word)
         n = len(encodings)
         ones = sum(encodings)
         if n == 0: