initial commit

.gitattributes

@@ -0,0 +1,31 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,13 @@
+---
+title: Winnie
+emoji: 🔥
+colorFrom: indigo
+colorTo: indigo
+sdk: gradio
+sdk_version: 3.1.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,25 @@
+
+import gradio as gr
+
+from bot.interface import Chatbot
+
+bot = Chatbot()
+
+def greet(input_txt, history = []):
+    global bot
+
+    if bot is None:
+        bot = Chatbot()
+
+    history.append(input_txt)
+    response = bot.chat(history)
+    history.append(response)
+
+    return response, history
+
+if __name__ == "__main__":
+    gr.Interface(fn=greet,
+        # title="使用中文和脑子瓦特了的Vicky聊天",
+        inputs=["text", "state"],
+        outputs=["text", "state"]
+    ).launch()

bot/interface.py

@@ -0,0 +1,48 @@
+
+from random import choice
+from random import randint
+from random import uniform
+
+from bot.simctgdialogue import SimCTGDialogue
+
+class Chatbot():
+    def __init__(self):
+        self.model = SimCTGDialogue("cambridgeltl/simctg_lccc_dialogue", [])
+        self.tokenizer = self.model.tokenizer
+        self.model.eval()
+
+    def __contrastive_search(self, context_list):
+        print("__contrastive_search")
+        print(context_list)
+        beam_width, alpha, decoding_len = randint(1, 8), uniform(0.10, 0.40), 64
+        return self.model.contrastive_search(context_list, beam_width, alpha, decoding_len)
+
+    def __diverse_contrastive_search(self, context_list):
+        print("__diverse_contrastive_search")
+        print(context_list)
+        sample_step, nucleus_p = 1, uniform(0.10, 0.40)
+        beam_width, alpha, decoding_len = randint(1, 5), uniform(0.10, 0.40), 64
+        return self.model.diverse_contrastive_search(context_list, sample_step, nucleus_p, beam_width, alpha, decoding_len)
+
+    def __greedy_search(self, context_list):
+        print("__greedy_search")
+        print(context_list)
+        decoding_len = 64
+        return self.model.greedy_search(context_list, decoding_len)
+
+    def __beam_search(self, context_list):
+        print("__beam_search")
+        print(context_list)
+        beam_width, decoding_len = randint(1, 9), 64
+        return self.model.beam_search(context_list, beam_width, decoding_len)
+
+    def chat(self, prefix = []):
+        methods_for_sort_dialogue = [self.__contrastive_search, self.__greedy_search]
+        methods_for_long_dialogue = [self.__beam_search, self.__diverse_contrastive_search, self.__greedy_search, self.__contrastive_search]
+
+        if ( len(prefix) < 4 ):
+            response = choice(methods_for_sort_dialogue)(prefix)
+        else:
+            response = choice(methods_for_long_dialogue)(prefix)
+
+        return response

bot/simctgdialogue.py

@@ -0,0 +1,177 @@
+
+import torch
+
+from torch import nn
+
+class SimCTGDialogue(nn.Module):
+    def __init__(self, model_name, additional_special_tokens):
+        super(SimCTGDialogue, self).__init__()
+        from transformers import AutoTokenizer, GPT2LMHeadModel
+        eos_token = '[SEP]'
+        pad_token = '[PAD]'
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name, additional_special_tokens=additional_special_tokens)
+        self.vocab_size = len(self.tokenizer)
+        self.model = GPT2LMHeadModel.from_pretrained(model_name)
+        self.embed_dim = self.model.config.hidden_size
+        if pad_token in self.tokenizer.vocab:
+            print ('PAD token exists.')
+        else:
+            print ('Add PAD token to the tokenizer.')
+            print ('Original vocabulary size is {}'.format(len(self.tokenizer)))
+            self.tokenizer.add_tokens([pad_token])
+            print ('Vocabulary size after extension is {}'.format(len(self.tokenizer)))
+            assert len(self.tokenizer.convert_tokens_to_ids([pad_token])) == 1
+            self.model.resize_token_embeddings(len(self.tokenizer)) 
+        self.pad_token_id = self.tokenizer.convert_tokens_to_ids([pad_token])[0]
+        self.vocab_size = len(self.tokenizer)
+        if 'e' in eos_token:
+            self.eos_token = self.tokenizer.eos_token
+        else:
+            self.eos_token = eos_token
+        print (self.eos_token)
+
+    def parse_dialogue_context(self, context_list, cuda_available=False, device=0):
+        # context_list: a list of utterances in the dialogue session
+        uttr_num = len(context_list)
+        context_text = self.eos_token.join(context_list).strip(self.eos_token) + self.eos_token
+        #print (context_text)
+        tokens = self.tokenizer.tokenize(context_text)
+        input_ids = self.tokenizer.convert_tokens_to_ids(tokens)
+        input_ids = input_ids
+        input_ids = torch.LongTensor(input_ids).view(1,-1)
+        if cuda_available:
+            input_ids = input_ids.cuda(device)
+        return input_ids, uttr_num
+
+    def extract_response(self, output_ids, uttr_num):
+        output_text = self.tokenizer.decode(output_ids)
+        # extract response
+        item_list = output_text.split(self.eos_token)
+        response = item_list[uttr_num].strip()
+        if self.eos_token == '<|endoftext|>': # English GPT
+            response = ' '.join(response.split())
+        else:
+            response = ''.join(response.split())
+        return response
+
+    def contrastive_search(self, context_list, beam_width, alpha, decoding_len, 
+        cuda_available=False, device=0):
+        input_ids, uttr_num = self.parse_dialogue_context(context_list, 
+            cuda_available=cuda_available, device=device)
+        output = self.fast_contrastive_generation(input_ids, beam_width, alpha, decoding_len)
+        return self.extract_response(output, uttr_num)
+
+    def diverse_contrastive_search(self, context_list, sample_step, nucleus_p, 
+        beam_width, alpha, decoding_len, cuda_available=False, device=0):
+        input_ids, uttr_num = self.parse_dialogue_context(context_list,
+            cuda_available=cuda_available, device=device)
+        output = self.diverse_contrastive_generation(input_ids, sample_step, nucleus_p, 
+            beam_width, alpha, decoding_len)
+        return self.extract_response(output, uttr_num)
+
+    def greedy_search(self, context_list, decoding_len, cuda_available=False, device=0):
+        input_ids, uttr_num = self.parse_dialogue_context(context_list,
+            cuda_available=cuda_available, device=device)
+        output = self.greedy_generation(input_ids, decoding_len)
+        return self.extract_response(output, uttr_num)
+
+    def beam_search(self, context_list, beam_width, decoding_len, 
+        cuda_available=False, device=0):
+        input_ids, uttr_num = self.parse_dialogue_context(context_list,
+            cuda_available=cuda_available, device=device)
+        output = self.beam_generation(input_ids, beam_width, decoding_len)
+        return self.extract_response(output, uttr_num)
+
+    def nucleus_sampling(self, context_list, nucleus_p, decoding_len, 
+        cuda_available=False, device=0):
+        input_ids, uttr_num = self.parse_dialogue_context(context_list,
+            cuda_available=cuda_available, device=device)
+        output = self.nucleus_generation(input_ids, nucleus_p, decoding_len)
+        return self.extract_response(output, uttr_num)
+
+    def fast_contrastive_generation(self, input_ids, beam_width, alpha, decoding_len):
+        '''
+           input_ids: prefix input; 1 x prefix_len
+           decoding_len: how many tokens to generate
+           beam_width: size of candidate pool during decoding
+           alpha: regulates importance of model confidence and degeneration penalty
+        '''
+        self.model.eval()
+        from bot.utlis import ContrastiveDecodingOneStepFast
+        # sanity check
+        assert alpha >= 0. and alpha <= 1.0
+        
+        # fast mode
+        batch_size, seqlen = input_ids.size()
+        #generated = [[] for _ in range(batch_size)]
+        generated = [item for item in input_ids.tolist()]
+        past_key_values = None
+        last_hidden_states = None
+        logits = None
+        for step in range(decoding_len):
+            input_ids, past_key_values, last_hidden_states, logits = ContrastiveDecodingOneStepFast(
+                self.model,
+                input_ids,
+                beam_width,
+                alpha,
+                past_key_values,
+                last_hidden_states,
+                self.tokenizer,
+                logits,
+                first_step=step == 0,
+            )
+            tokens = input_ids.squeeze(dim=-1).tolist()
+            for idx, t in enumerate(tokens):
+                generated[idx].append(t)
+        return generated[0]
+
+    def diverse_contrastive_generation(self, input_ids, sample_step, nucleus_p, beam_width, alpha, decoding_len):
+        '''
+            sample_step: 
+                number of steps to decode with nucleus sampling, 
+                for the remaining steps we use contrastive search
+            decoding_len: 
+                the total number of generated tokens
+            beam_width: 
+                size of candidate pool during decoding
+            alpha: 
+                regulates importance of model confidence and degeneration penalty
+
+        '''
+        contrastive_step = decoding_len - sample_step
+        _, prefix_len = input_ids.size()
+        # first do sample
+        input_ids = self.model.generate(
+                            input_ids, 
+                            do_sample=True, 
+                            max_length=prefix_len+sample_step, 
+                            top_p=nucleus_p,
+                            top_k=0)
+        # then do contrastive search
+        output = self.fast_contrastive_generation(input_ids, beam_width, alpha, contrastive_step)
+        return output
+
+    def greedy_generation(self, input_ids, decoding_len):
+        _, prefix_len = input_ids.size()
+        output = self.model.generate(
+                            input_ids, 
+                            max_length=prefix_len+decoding_len)
+        return output[0]
+
+    def beam_generation(self, input_ids, beam_width, decoding_len):
+        _, prefix_len = input_ids.size()
+        output = self.model.generate(
+                            input_ids, 
+                            max_length=prefix_len+decoding_len, 
+                            num_beams=beam_width)
+        return output[0]
+
+    def nucleus_generation(self, input_ids, nucleus_p, decoding_len):
+        _, prefix_len = input_ids.size()
+        output = self.model.generate(
+                            input_ids, 
+                            do_sample=True, 
+                            max_length=prefix_len+decoding_len, 
+                            top_p=nucleus_p,
+                            top_k=0)
+        return output[0]

bot/utlis.py

@@ -0,0 +1,174 @@
+
+import torch
+import random
+import torch.nn.functional as F
+
+def ranking(context_hidden, next_hidden, next_top_k_ids, next_top_k_probs, alpha):
+    '''
+        context_hidden: beam_width x context_len x embed_dim
+        next_hidden: beam_width x 1 x embed_dim
+        next_top_k_ids: beam_width x 1
+    '''
+    beam_width, context_len, embed_dim = context_hidden.size()
+    assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])
+    norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True)
+    norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True)
+    cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1,2)).squeeze(-1)
+    assert cosine_matrix.size() == torch.Size([beam_width, context_len])
+    scores, _ = torch.max(cosine_matrix, dim = -1)
+    assert scores.size() == torch.Size([beam_width])
+    next_top_k_probs = next_top_k_probs.view(-1)
+    scores = (1.0 - alpha) * next_top_k_probs - alpha * scores 
+    _, selected_idx = torch.topk(scores, k = 1)
+    assert selected_idx.size() == torch.Size([1])
+    selected_idx = selected_idx.unsqueeze(0)
+    assert selected_idx.size() == torch.Size([1,1])
+    next_id = torch.gather(next_top_k_ids, dim = 0, index=selected_idx)
+    assert next_id.size() == torch.Size([1,1])
+    return next_id
+
+def ContrastiveDecodingOneStep(model, input_ids, beam_width, alpha):
+    '''
+        model: the generation model, e.g., gpt2
+        input_ids: 1 x seqlen
+    '''
+    prev_hidden_states, logits = model.compute_logits_and_hidden_states(input_ids)
+    _, seqlen, embed_dim = prev_hidden_states.size()
+    _, _, vocab_size = logits.size()
+    p = random.uniform(0, 1)
+
+    logit_for_next_step = logits[:,-1,:]
+    assert logit_for_next_step.size() == torch.Size([1, vocab_size])
+
+    next_probs = F.softmax(logit_for_next_step, dim = -1)
+    assert next_probs.size() == logit_for_next_step.size()
+
+    _, top_k_ids = torch.topk(logit_for_next_step, dim = -1, k = beam_width)
+    assert top_k_ids.size() == torch.Size([1, beam_width])
+        
+    top_k_probs = torch.gather(next_probs, dim = 1, index=top_k_ids)
+
+    assert top_k_probs.size() == top_k_ids.size()
+    # compute new hidden 
+    expanded_context = [input_ids for _ in range(beam_width)]
+    expanded_context = torch.cat(expanded_context, dim = 0)
+    assert expanded_context.size() == torch.Size([beam_width, seqlen])
+    top_k_ids = top_k_ids.view(beam_width, 1)
+    next_input_ids = torch.cat([expanded_context, top_k_ids], dim = -1)
+    assert next_input_ids.size() == torch.Size([beam_width, seqlen+1])
+    new_hidden_states, next_logits = model.compute_logits_and_hidden_states(next_input_ids)
+    assert new_hidden_states.size() == torch.Size([beam_width, seqlen+1, embed_dim])
+    context_hidden = new_hidden_states[:,:seqlen,:]
+    assert context_hidden.size() == torch.Size([beam_width, seqlen, embed_dim])
+    next_hidden = new_hidden_states[:,seqlen:,:]
+    assert next_hidden.size() == torch.Size([beam_width, 1, embed_dim])
+
+    next_id = ranking(context_hidden, next_hidden, top_k_ids, top_k_probs, alpha)       
+
+    next_input_ids = torch.cat([input_ids, next_id], dim = -1)
+    assert next_input_ids.size() == torch.Size([1, seqlen+1])
+    return next_input_ids
+
+# ========== batch version ========= #
+def ranking_fast(context_hidden, next_hidden, next_top_k_probs, alpha, beam_width):
+    '''
+        context_hidden: bsz*beam x seqlen x embed_dim
+        next_hidden: bsz*beam x 1 x embed_dim
+        next_top_k_probs: bsz x beam
+    '''
+    _, context_len, embed_dim = context_hidden.size()
+    norm_context_hidden = context_hidden / context_hidden.norm(dim=2, keepdim=True)
+    norm_next_hidden = next_hidden / next_hidden.norm(dim=2, keepdim=True)
+    cosine_matrix = torch.matmul(norm_context_hidden, norm_next_hidden.transpose(1,2)).squeeze(-1)    # [B*K, S]
+    scores, _ = torch.max(cosine_matrix, dim=-1)    # [B*K]
+    next_top_k_probs = next_top_k_probs.view(-1)    # [B*K]
+    scores = (1.0 - alpha) * next_top_k_probs - alpha * scores 
+    scores = torch.stack(torch.split(scores, beam_width))    # [B, K]
+    selected_idx = scores.max(dim=-1)[1]    # [B]
+    return selected_idx
+
+def ContrastiveDecodingOneStepFast(
+    model, 
+    ids, 
+    beam_width, 
+    alpha, 
+    past_key_values,
+    last_hidden_states,
+    vocab,
+    logit_for_next_step,
+    first_step=False,
+    ):
+    # input_ids: [B, S]
+    if first_step:
+        output = model(
+            input_ids=ids, 
+            past_key_values=past_key_values,
+            use_cache=True,
+            output_hidden_states=True
+        )
+        past_key_values = output.past_key_values
+        last_hidden_states = output.hidden_states[-1]    # [B, S, E]
+        logit_for_next_step = output.logits[:, -1, :]    # [B, V]
+    bsz, seqlen, embed_dim = last_hidden_states.size()
+    p = random.uniform(0, 1)
+
+    next_probs = F.softmax(logit_for_next_step, dim=-1)
+    _, top_k_ids = torch.topk(logit_for_next_step, dim=-1, k=beam_width)    # [B, K]
+    top_k_probs = torch.gather(next_probs, dim=1, index=top_k_ids)    # [B, K]
+    # compute new hidden
+    past_key_values = enlarge_past_key_values(past_key_values, beam_width)
+    output = model(
+        input_ids=top_k_ids.view(-1, 1), 
+        attention_mask=torch.ones_like(top_k_ids.view(-1, 1)),
+        past_key_values=past_key_values,
+        output_hidden_states=True,
+        use_cache=True,
+    )
+    past_key_values = output.past_key_values
+    logits = output.logits[:, -1, :]    # [B*K, V]
+    next_hidden = output.hidden_states[-1]    # [B*K, 1, E]
+    context_hidden = last_hidden_states.unsqueeze(1).expand(-1, beam_width, -1, -1).reshape(bsz*beam_width, seqlen, embed_dim)    # [B*K, S, E]
+
+    selected_idx = ranking_fast(
+        context_hidden, 
+        next_hidden, 
+        top_k_probs,    # [B, K] 
+        alpha,
+        beam_width,
+    )     # [B]
+    # prepare for the next step
+    next_id = top_k_ids[range(len(top_k_ids)), selected_idx].unsqueeze(-1)    # [B, 1]
+    next_hidden = torch.stack(torch.split(next_hidden.squeeze(dim=1), beam_width))    # [B, K, E]
+    next_hidden = next_hidden[range(bsz), selected_idx, :]    # [B, E]
+    last_hidden_states = torch.cat([last_hidden_states, next_hidden.unsqueeze(1)], dim=1)    # [B, S, E]
+    past_key_values = select_past_key_values(past_key_values, beam_width, selected_idx)
+    logits = torch.stack(torch.split(logits, beam_width))[range(bsz), selected_idx, :]    # [B, V]
+    # next_id: [B, 1]
+    return next_id, past_key_values, last_hidden_states, logits 
+
+def enlarge_past_key_values(past_key_values, beam_width):
+    # from [B, num_head, seq_len, esz] to [B*K, num_head, seq_len, esz]
+    new_key_values = []
+    for layer in past_key_values:
+        items = []
+        for item in layer:
+            # item is the key and value matrix
+            bsz, num_head, seq_len, esz = item.size()
+            item = item.unsqueeze(1).expand(-1, beam_width, -1, -1, -1).reshape(bsz*beam_width, num_head, seq_len, esz)    # [bsz*beam, num_head, seq_len, esz]
+            items.append(item)
+        new_key_values.append(items)
+    return new_key_values
+
+def select_past_key_values(past_key_values, beam_width, selected_idx):
+    '''select_idx: [B]'''
+    new_key_values = []
+    for layer in past_key_values:
+        items = []
+        for item in layer:
+            bsz_and_beam, num_head, seq_len, esz = item.size()
+            bsz = int(bsz_and_beam//beam_width)
+            item = torch.stack(torch.split(item, beam_width, dim=0))    # [B, K, num_head, seq_len, esz] 
+            item = item[range(bsz), selected_idx, :, :, :]   # [B, num_head, seq_len, esz]
+            items.append(item)
+        new_key_values.append(items)
+    return new_key_values

requirements.txt

@@ -0,0 +1,18 @@
+absl-py
+pytest
+sacrebleu==1.4.10
+six
+wheel
+progressbar
+sklearn
+torch==1.6.0 
+torchvision==0.7.0
+transformers==4.7.0
+pyyaml
+nltk
+sentencepiece
+spacy
+gdown
+seaborn
+matplotlib
+pandas