first commit, add gitignore

.gitignore

@@ -0,0 +1,3 @@
+/__pycache__
+
+.pyc

app.py

@@ -15,29 +15,7 @@ def respond(
     temperature,
     top_p,
 ):
-    messages = [{"role": "system", "content": system_message}]
-
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
-
-        response += token
-        yield response
+    return "test", []
 
 """
 For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
@@ -45,16 +23,9 @@ For information on how to customize the ChatInterface, peruse the gradio docs: h
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
+        gr.Markdown("Please enter your message. Add privacy tags (<sensitive>...</sensitive>) around the words you want to hide. Only the most recent message submitted will be taken into account (no history is retained)."),
+        gr.Slider(minimum=1, maximum=128, value=32, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
     ],
 )
 

models/modeling_llamask.py

@@ -0,0 +1,117 @@
+# coding=utf-8
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# 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.
+
+from typing import List, Optional, Tuple, Union
+
+import torch
+from torch import nn
+
+from transformers.cache_utils import Cache
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.llama import LlamaForCausalLM
+
+class LlamaskForCausalLM(LlamaForCausalLM):
+    def __init__(self, config):
+        super().__init__(config)
+        self.special_tokens = nn.Embedding(2, config.hidden_size) # 0 -> mask encoding, 1 -> buffer token
+        self.post_init()
+
+    def generate(
+            self,
+            input_ids: torch.LongTensor = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            max_tokens: int=32,
+            temperature: float=1.0,
+    ):
+        eos_token_tensor = torch.tensor(self.config.eos_token_id, device=input_ids.device)
+        for _ in range(max_tokens):
+            outputs = self.forward(input_ids=input_ids, attention_mask=attention_mask)
+            logits = outputs['logits'][:,-1,:]/temperature
+
+            probs = torch.nn.functional.softmax(logits, dim=-1)
+            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+            batch_size, seq_len, _ = attention_mask.shape
+            expanded_mask = torch.zeros(batch_size, seq_len + 1, seq_len + 1, dtype=attention_mask.dtype, device=attention_mask.device)
+
+            # Step 1: Copy the existing attention mask (top-left block of the expanded mask)
+            expanded_mask[:, :seq_len, :seq_len] = attention_mask
+
+            # Step 2: Copy the last row of the original attention mask into the new row (excluding the last position)
+            expanded_mask[:, seq_len, :seq_len] = attention_mask[:, -1, :]
+
+            # Step 3: Set the diagonal of the new token to attend to all previous tokens by setting the new last element to 1
+            expanded_mask[:, seq_len, seq_len] = 1
+
+            next_tokens = next_tokens[:, None]
+
+            input_ids = torch.cat([input_ids, next_tokens], dim=-1)
+            attention_mask = expanded_mask
+            if torch.all(torch.any(next_tokens==eos_token_tensor, dim=1)):
+                break
+        return input_ids
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor = None,
+        num_buffer_token: Optional[int] = 0,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+        batch_size = input_ids.shape[0]
+
+        # print("BEWARE PRIVACY TAG DISABLE")
+        # privacy_tag = self.special_tokens(torch.tensor([0], device=input_ids.device))
+        # buffer_token = self.special_tokens(torch.tensor([0], device=input_ids.device)).unsqueeze(0)
+        
+        inputs_embeds = self.model.embed_tokens(input_ids)
+        # buffer_tokens = buffer_token.repeat(batch_size, num_buffer_token, 1)
+        # inputs_embeds = torch.cat([inputs_embeds, buffer_tokens], dim=1)
+        
+        # inputs_embeds[attention_mask[:,-1,:]==0] = inputs_embeds[attention_mask[:,-1,:]==0] + privacy_tag
+
+        attention_mask = attention_mask.unsqueeze(1)
+        attention_mask = attention_mask.to(inputs_embeds.dtype)
+        attention_mask = attention_mask.masked_fill(attention_mask == 0, -1e9)
+        attention_mask = attention_mask.masked_fill(attention_mask == 1, float(0.0))
+        
+        outputs = super().forward(
+                input_ids=None,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                return_dict=return_dict,
+                cache_position=cache_position,
+            )
+        return outputs
+    

models/tokenizer_utils.py

@@ -0,0 +1,83 @@
+from typing import List
+
+from transformers import AutoTokenizer
+import torch
+
+def prepare_tokenizer(tokenizer: AutoTokenizer, token_beg="<sensitive>", token_end="</sensitive>"):
+    """
+    Add privacy special tokens
+    """
+    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer.add_special_tokens({"additional_special_tokens": [token_beg, token_end]})
+    tokenizer.sensitive_beg_id = tokenizer.encode(token_beg, add_special_tokens=False)[0]
+    tokenizer.sensitive_end_id = tokenizer.encode(token_end, add_special_tokens=False)[0]
+
+def generate_custom_mask(tokenizer: AutoTokenizer, prompts: List[str], device='cpu', padding_side='left'):
+    """
+    Given a prepared tokenizer (i.e. with privacy special tokens), a list of prompts with privacy special tokens,
+    tokenize and generate custom masks for a privacy-compatible transformer.
+    """
+    input_ids = tokenizer(prompts)['input_ids']
+    return generate_custom_mask_input_ids(tokenizer, input_ids, device=device, padding_side='left')[0]
+
+def generate_custom_mask_input_ids(tokenizer: AutoTokenizer, input_ids, device='cpu', padding_side="right"):
+    """
+    Given a prepared tokenizer (i.e. with privacy special tokens), a list of prompts with privacy special tokens,
+    tokenize and generate custom masks for a privacy-compatible transformer.
+    """
+    new_input_ids, new_attention_masks, seq_len_list = [], [], []
+    max_len = 0
+    batch_size = len(input_ids)
+    for input_id in input_ids:
+        trigger_privacy = False
+        new_input_id = []
+        mask_pos_list = []
+        idx = 0
+        for token_id in input_id:
+            if token_id == tokenizer.sensitive_beg_id:
+                trigger_privacy = True
+            elif token_id == tokenizer.sensitive_end_id:
+                trigger_privacy = False
+            else:
+                new_input_id.append(token_id)
+                if trigger_privacy:
+                    mask_pos_list.append(idx)
+                idx += 1
+        seq_len = len(new_input_id)
+        seq_len_list.append(seq_len)
+
+        attention_mask = torch.tril(torch.ones((seq_len, seq_len)))
+        
+        for idx in mask_pos_list:
+            # The last token can access everything.
+            attention_mask[idx+1:-1, idx] = 0
+            attention_mask[idx,:idx] = 1
+        new_attention_masks.append(attention_mask)
+        new_input_ids.append(new_input_id)
+
+        max_len = max(max_len, seq_len)
+    
+    new_full_attention_mask = torch.zeros((batch_size, max_len))
+    for batch, seq_len in enumerate(seq_len_list):
+        if padding_side == 'left':
+            new_full_attention_mask[batch, max_len-seq_len:] = 1
+        else:
+            new_full_attention_mask[batch, :seq_len] = 1
+
+    for idx, (input_ids, attention_mask) in enumerate(zip(new_input_ids, new_attention_masks)):
+        current_len = len(input_ids)
+        new_attention_mask = torch.zeros((max_len, max_len), dtype=torch.long)
+        if padding_side == 'left':
+            input_ids = [tokenizer.pad_token_id]*(max_len - current_len) + input_ids
+        else:
+            input_ids = input_ids + [tokenizer.pad_token_id]*(max_len - current_len)
+        
+        if padding_side == 'left':
+            new_attention_mask[max_len-current_len:, max_len-current_len:] = attention_mask
+        else:
+            new_attention_mask[:current_len,:current_len] = attention_mask
+        new_input_ids[idx] = torch.tensor(input_ids).unsqueeze(0)
+        new_attention_masks[idx] = new_attention_mask.unsqueeze(0)
+    input_id = torch.cat(new_input_ids, dim=0)
+    attention_mask = torch.cat(new_attention_masks, dim=0)
+    return {'input_ids': input_id.to(device), 'attention_mask': attention_mask.to(device)}, new_full_attention_mask.to(device)