Refactoring

.gitattributes

@@ -9,13 +9,9 @@
 *.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

README.md

@@ -1,10 +1,10 @@
 ---
 title: Winnie
-emoji: 🔥
-colorFrom: indigo
-colorTo: indigo
+emoji: ⚡
+colorFrom: pink
+colorTo: gray
 sdk: gradio
-sdk_version: 3.1.1
+sdk_version: 3.0.24
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,25 +1,78 @@
 
+import os
+
 import gradio as gr
 
-from bot.interface import Chatbot
+from bot.chatbot import ChatBot
+from bot.config import special_token_list
+
+bot:ChatBot = None
 
-bot = Chatbot()
+def get_skill_list() -> list:
+    path:str = os.path.split( os.path.realpath(__file__) )[0]
+    file_list:list[str] = os.listdir( path + "/bot/skills/" )
+    plugin_list:list[str] = []
+    for file in file_list:
+        if file.endswith(".py"):
+            plugin_list.append( file[:-3] )
+    return plugin_list
 
-def greet(input_txt, history = []):
+def general(input_txt:str, state:dict = {}):
     global bot
 
-    if bot is None:
-        bot = Chatbot()
+    history_list:list = state.get("history", [])
+    role_card:dict[str, str] = state.get("role_card", {
+        "<NAME>": "Winnie",
+        "<GENDER>": "女",
+        "<YEAROFBIRTH>":"1995",
+        "<MONTHOFBIRTH>":"5",
+        "<DAYOFBIRTH>":"6",
+        "<ZODIAC>":"金牛座",
+        "<AGE>":"27"
+        }
+    )
 
-    history.append(input_txt)
-    response = bot.chat(history)
-    history.append(response)
+    output_txt:str = None
 
-    return response, history
+    for skill_name in get_skill_list():
+        if output_txt is None:
+            plugin = __import__("bot.skills."+skill_name, fromlist=[skill_name])
+            plugin_class = getattr(plugin, "Skill")
+            p = plugin_class()
+            output_txt, history_list, role_card = p.process(input_txt, history_list, role_card)
 
-if __name__ == "__main__":
-    gr.Interface(fn=greet,
-        # title="使用中文和脑子瓦特了的Vicky聊天",
-        inputs=["text", "state"],
-        outputs=["text", "state"]
+    if output_txt is None:
+        res, history_list = bot.chat(input_txt, history_list, role_card=role_card)
+        output_txt = "".join(res)
+
+    state["history"] = history_list
+    state["role_card"] = role_card
+
+    return output_txt, state
+
+def main() -> None:
+    global bot
+
+    bot = ChatBot.get_chat_bot("lewiswu1209/Winnie", special_token_list=special_token_list)
+
+    title:str = "使用中文和Winnie聊天"
+
+    description:str = "输入任意文字，Winnie会和你对话<br>"
+    description += "输入ERASE MEMORY，会清空Winnie的记忆<br>"
+    description += "输入\"&lt;TAG&gt;=&lt;VALUE&gt;\"，可以修改Winnie的角色信息<br>"
+    description += "例如：&lt;NAME&gt;=Vicky，会修改Winnie的名字<br>"
+    description += "可以修改的角色信息有：<br>"
+    description += "&lt;NAME&gt;, &lt;GENDER&gt;, &lt;YEAROFBIRTH&gt;, &lt;MONTHOFBIRTH&gt;, &lt;DAYOFBIRTH&gt;, &lt;ZODIAC&gt;, &lt;AGE&gt;<br>"
+    description += "输入“上联：XXXXXXX”，Winnie会和你对对联<br>"
+    description += "输入“写诗：XXXXXXX”，Winnie会以XXXXXXX为开头写诗"
+
+    gr.Interface(
+        fn = general,
+        title = title,
+        description = description,
+        inputs = ["text", "state"],
+        outputs = ["text", "state"]
     ).launch()
+
+if __name__ == "__main__":
+    main()

bot/chatbot.py

@@ -0,0 +1,109 @@
+
+from tokenize import tokenize
+import torch
+import torch.nn.functional as F
+
+from transformers import BertTokenizer, GPT2LMHeadModel
+
+def replace_special_tokens(tokens:list, special_token_dict:dict)->list:
+    replaced_tokens:list = []
+
+    for token in tokens:
+        if token in special_token_dict.keys():
+            replaced_tokens.append( special_token_dict[token] )
+        else:
+            replaced_tokens.append( token )
+
+    return replaced_tokens
+
+def top_k_top_p_filtering(logits, top_k:int=0, top_p:float=0.0, filter_value:float=-float('Inf')):
+    top_k = min( top_k, logits.size(-1) )
+
+    if top_k > 0:
+        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+        logits[indices_to_remove] = filter_value
+
+    if top_p > 0.0:
+        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+        sorted_indices_to_remove[..., 0] = 0
+        indices_to_remove = sorted_indices[sorted_indices_to_remove]
+        logits[indices_to_remove] = filter_value
+
+    return logits
+
+class ChatBot():
+
+    def get_chat_bot(checkpoint:str, vocab_path:str = None, special_token_list:list = [])->object:
+        tokenizer = ChatBot.get_tokenizer(checkpoint, vocab_path, special_token_list)
+        model = GPT2LMHeadModel.from_pretrained(checkpoint)
+
+        return ChatBot(tokenizer, model)
+
+    def get_tokenizer(checkpoint:str, vocab_path:str = None, special_token_list:list = [])->object:
+        if vocab_path is None:
+            tokenizer = BertTokenizer.from_pretrained(checkpoint, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
+        else:
+            tokenizer = BertTokenizer(vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
+        tokenizer.add_special_tokens( {'additional_special_tokens':special_token_list} )
+
+        return tokenizer
+
+    def __init__(self, tokenizer:object, model:object)->None:
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.tokenizer = tokenizer
+        self.model = model
+        self.model.to(self.device)
+
+    def convert_ids_to_tokens(self, ids:list):
+        return self.tokenizer.convert_ids_to_tokens(ids)
+
+    def convert_ids_to_text(self, ids):
+        return "".join( self.convert_ids_to_tokens(ids) )
+
+    def convert_text_to_ids(self, text, add_special_tokens=False):
+        return self.tokenizer.encode(text, add_special_tokens=add_special_tokens)
+
+    def get_prediction(self, input_tensor, input_ids, repetition_penalty, temperature, top_k, top_p):
+        self.model.eval()
+
+        generated_ids = []
+        for _ in range(64):
+            output_pt = self.model(input_tensor)
+
+            next_token_logits = output_pt.logits[0, -1, :]
+            for id in set(input_ids):
+                if id != self.tokenizer.sep_token_id:
+                    next_token_logits[id] /= repetition_penalty
+            for id in set(generated_ids):
+                next_token_logits[id] /= repetition_penalty
+            next_token_logits = next_token_logits / temperature
+            next_token_logits[self.tokenizer.convert_tokens_to_ids('[UNK]')] = -float('Inf')
+            filtered_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
+            next_token = torch.multinomial( F.softmax(filtered_logits, dim=-1), num_samples=1 )
+
+            if next_token == self.tokenizer.sep_token_id:
+                break
+
+            input_tensor = torch.cat( (input_tensor, next_token.unsqueeze(0)), dim=1 )
+            generated_ids.append( next_token.item() )
+
+        return generated_ids
+
+    def chat(self:object, text:str, history:list, role_card:dict={}) -> str:
+        text_ids = self.tokenizer.encode(text, add_special_tokens=False)
+        history.append(text_ids)
+        input_ids = [self.tokenizer.cls_token_id]
+        for history_utr in history[-50:]:
+            input_ids.extend(history_utr)
+            input_ids.append(self.tokenizer.sep_token_id)
+        input_tensor = torch.tensor(input_ids).to(self.device).unsqueeze(0)
+        generated_ids = self.get_prediction(input_tensor, input_ids, repetition_penalty=1.2, temperature=0.73, top_k=10, top_p=0.7)
+
+        history.append(generated_ids)
+
+        generated_tokens = replace_special_tokens( self.convert_ids_to_tokens(generated_ids), role_card )
+
+        return "".join(generated_tokens), history

bot/config.py

@@ -0,0 +1,21 @@
+special_token_list:list = [
+        '<NAME>',
+        '<GENDER>',
+        '<YEAROFBIRTH>',
+        '<MONTHOFBIRTH>',
+        '<DAYOFBIRTH>',
+        '<ZODIAC>',
+        '<AGE>',
+        '<CMD>',
+        '<NICK>',
+        '<HEIGHT>',
+        '<WEIGHT>',
+        '<WORK>',
+        '<HOBBY>',
+        '<HOMETOWN>',
+        '<CITY>',
+        '<BUST>',
+        '<WAIST>',
+        '<HIP>',
+        '<CUP>'
+    ]

bot/interface.py

@@ -1,48 +0,0 @@
-
-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

@@ -1,177 +0,0 @@
-
-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/skills/couplet.py

@@ -0,0 +1,16 @@
+
+import requests
+
+class Skill:
+    def __init__(self:object) -> None:
+        pass
+
+    def process(self:object, input_txt:str, history_list:list, role_card:dict):
+        output_text:str = None
+        if input_txt.startswith("上联：") or input_txt.startswith("上联:"):
+            output_text = requests.post(
+                url='https://hf.space/embed/lewiswu1209/gpt2-chinese-couplet/+/api/predict/',
+                json={"data": [input_txt[3:]]}
+            ).json()["data"][0]
+            output_text = "我对下联：" + output_text
+        return  output_text, history_list, role_card

bot/skills/delete_memory.py

@@ -0,0 +1,11 @@
+
+class Skill:
+    def __init__(self:object) -> None:
+        pass
+
+    def process(self:object, input_txt:str, history_list:list, role_card:dict):
+        output_txt:str = None
+        if input_txt.upper()=="ERASE MEMORY":
+            history_list = []
+            output_txt = "我是谁？我在哪？我在干什么？"
+        return  output_txt, history_list, role_card

bot/skills/give_role.py

@@ -0,0 +1,14 @@
+
+class Skill:
+    def __init__(self:object) -> None:
+        pass
+
+    def process(self:object, input_txt:str, history_list:list, role_card:dict):
+        output_txt:str = None
+        for tag in role_card.keys():
+            prefix:str = "{}=".format(tag)
+            if input_txt.startswith( prefix ):
+                role_card[tag]=input_txt[len(prefix):]
+                output_txt = "已设置{}为{}".format(tag, role_card[tag])
+                break
+        return  output_txt, history_list, role_card

bot/skills/poem.py

@@ -0,0 +1,15 @@
+
+import requests
+
+class Skill:
+    def __init__(self:object) -> None:
+        pass
+
+    def process(self:object, input_txt:str, history_list:list, role_card:dict):
+        output_text:str = None
+        if input_txt.startswith("写诗：") or input_txt.startswith("写诗:"):
+            output_text = requests.post(
+                url='https://hf.space/embed/lewiswu1209/gpt2-chinese-poem/+/api/predict/',
+                json={"data": [input_txt[3:]]}
+            ).json()["data"][0]
+        return  output_text, history_list, role_card

bot/utlis.py

@@ -1,174 +0,0 @@
-
-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

data_parallel.py

@@ -0,0 +1,100 @@
+from torch.nn.parallel import DataParallel
+import torch
+from torch.nn.parallel._functions import Scatter
+from torch.nn.parallel.parallel_apply import parallel_apply
+
+
+def scatter(inputs, target_gpus, chunk_sizes, dim=0):
+    r"""
+    Slices tensors into approximately equal chunks and
+    distributes them across given GPUs. Duplicates
+    references to objects that are not tensors.
+    """
+    def scatter_map(obj):
+        if isinstance(obj, torch.Tensor):
+            try:
+                return Scatter.apply(target_gpus, chunk_sizes, dim, obj)
+            except:
+                print('obj', obj.size())
+                print('dim', dim)
+                print('chunk_sizes', chunk_sizes)
+                quit()
+        if isinstance(obj, tuple) and len(obj) > 0:
+            return list(zip(*map(scatter_map, obj)))
+        if isinstance(obj, list) and len(obj) > 0:
+            return list(map(list, zip(*map(scatter_map, obj))))
+        if isinstance(obj, dict) and len(obj) > 0:
+            return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
+        return [obj for targets in target_gpus]
+
+    # After scatter_map is called, a scatter_map cell will exist. This cell
+    # has a reference to the actual function scatter_map, which has references
+    # to a closure that has a reference to the scatter_map cell (because the
+    # fn is recursive). To avoid this reference cycle, we set the function to
+    # None, clearing the cell
+    try:
+        return scatter_map(inputs)
+    finally:
+        scatter_map = None
+
+
+def scatter_kwargs(inputs, kwargs, target_gpus, chunk_sizes, dim=0):
+    r"""Scatter with support for kwargs dictionary"""
+    inputs = scatter(inputs, target_gpus, chunk_sizes, dim) if inputs else []
+    kwargs = scatter(kwargs, target_gpus, chunk_sizes, dim) if kwargs else []
+    if len(inputs) < len(kwargs):
+        inputs.extend([() for _ in range(len(kwargs) - len(inputs))])
+    elif len(kwargs) < len(inputs):
+        kwargs.extend([{} for _ in range(len(inputs) - len(kwargs))])
+    inputs = tuple(inputs)
+    kwargs = tuple(kwargs)
+    return inputs, kwargs
+
+
+class BalancedDataParallel(DataParallel):
+    def __init__(self, gpu0_bsz, *args, **kwargs):
+        self.gpu0_bsz = gpu0_bsz
+        super().__init__(*args, **kwargs)
+
+    def forward(self, *inputs, **kwargs):
+        if not self.device_ids:
+            return self.module(*inputs, **kwargs)
+        if self.gpu0_bsz == 0:
+            device_ids = self.device_ids[1:]
+        else:
+            device_ids = self.device_ids
+        inputs, kwargs = self.scatter(inputs, kwargs, device_ids)
+        # print('len(inputs)1: ', str(len(inputs)))
+        # print('self.device_ids[:len(inputs)]', str(self.device_ids[:len(inputs)]))
+        if len(self.device_ids) == 1:
+            return self.module(*inputs[0], **kwargs[0])
+        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
+        if self.gpu0_bsz == 0:
+            replicas = replicas[1:]
+        outputs = self.parallel_apply(replicas, device_ids, inputs, kwargs)
+        return self.gather(outputs, self.output_device)
+
+    def parallel_apply(self, replicas, device_ids, inputs, kwargs):
+        return parallel_apply(replicas, inputs, kwargs, device_ids[:len(inputs)])
+
+    def scatter(self, inputs, kwargs, device_ids):
+        bsz = inputs[0].size(self.dim)
+        num_dev = len(self.device_ids)
+        gpu0_bsz = self.gpu0_bsz
+        bsz_unit = (bsz - gpu0_bsz) // (num_dev - 1)
+        if gpu0_bsz < bsz_unit:
+            chunk_sizes = [gpu0_bsz] + [bsz_unit] * (num_dev - 1)
+            delta = bsz - sum(chunk_sizes)
+            for i in range(delta):
+                chunk_sizes[i + 1] += 1
+            if gpu0_bsz == 0:
+                chunk_sizes = chunk_sizes[1:]
+        else:
+            return super().scatter(inputs, kwargs, device_ids)
+
+        # print('bsz: ', bsz)
+        # print('num_dev: ', num_dev)
+        # print('gpu0_bsz: ', gpu0_bsz)
+        # print('bsz_unit: ', bsz_unit)
+        # print('chunk_sizes: ', chunk_sizes)
+        return scatter_kwargs(inputs, kwargs, device_ids, chunk_sizes, dim=self.dim)

dataset.py

@@ -0,0 +1,21 @@
+from torch.utils.data import Dataset
+import torch
+
+
+class MyDataset(Dataset):
+    """
+
+    """
+
+    def __init__(self, input_list, max_len):
+        self.input_list = input_list
+        self.max_len = max_len
+
+    def __getitem__(self, index):
+        input_ids = self.input_list[index]
+        input_ids = input_ids[:self.max_len]
+        input_ids = torch.tensor(input_ids, dtype=torch.long)
+        return input_ids
+
+    def __len__(self):
+        return len(self.input_list)

preprocess.py

@@ -0,0 +1,105 @@
+from tokenizers import BertWordPieceTokenizer
+from transformers import BertTokenizer
+from transformers import BertTokenizerFast
+import argparse
+import pandas as pd
+import pickle
+import jieba.analyse
+from tqdm import tqdm
+from transformers import GPT2TokenizerFast, GPT2LMHeadModel
+import logging
+import numpy as np
+from chatbot.config import config
+
+
+def create_logger(log_path):
+    """
+    将日志输出到日志文件和控制台
+    """
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+
+    formatter = logging.Formatter(
+        '%(asctime)s - %(levelname)s - %(message)s')
+
+    # 创建一个handler，用于写入日志文件
+    file_handler = logging.FileHandler(
+        filename=log_path)
+    file_handler.setFormatter(formatter)
+    file_handler.setLevel(logging.INFO)
+    logger.addHandler(file_handler)
+
+    # 创建一个handler，用于将日志输出到控制台
+    console = logging.StreamHandler()
+    console.setLevel(logging.DEBUG)
+    console.setFormatter(formatter)
+    logger.addHandler(console)
+
+    return logger
+
+
+def preprocess():
+    """
+    对原始语料进行tokenize，将每段对话处理成如下形式："[CLS]utterance1[SEP]utterance2[SEP]utterance3[SEP]"
+    """
+    # 设置参数
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--vocab_path', default='vocab/vocab.txt', type=str, required=False,
+                        help='词表路径')
+    parser.add_argument('--log_path', default='data/preprocess.log', type=str, required=False, help='训练日志存放位置')
+    parser.add_argument('--train_path', default='data/train.txt', type=str, required=False, help='训练日志存放位置')
+    parser.add_argument('--save_path', default='data/train.pkl', type=str, required=False, help='tokenize的训练数据集')
+    args = parser.parse_args()
+
+    # 初始化日志对象
+    logger = create_logger(args.log_path)
+
+    # 初始化tokenizer
+    tokenizer = BertTokenizerFast(vocab_file=args.vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
+    special_tokens = []
+    for key in config["mask_token"].keys():
+        special_tokens.append(key)
+    tokenizer.add_special_tokens( {'additional_special_tokens':special_tokens} )
+    sep_id = tokenizer.sep_token_id
+    cls_id = tokenizer.cls_token_id
+    logger.info("preprocessing data,data path:{}, save path:{}".format(args.train_path, args.save_path))
+
+    # 读取训练数据集
+    with open(args.train_path, 'rb') as f:
+        data = f.read().decode("utf-8")
+
+    # 需要区分linux和windows环境下的换行符
+    if "\r\n" in data:
+        train_data = data.split("\r\n\r\n")
+    else:
+        train_data = data.split("\n\n")
+    logger.info("there are {} dialogue in dataset".format(len(train_data)))
+
+    # 开始进行tokenize
+    # 保存所有的对话数据,每条数据的格式为："[CLS]utterance1[SEP]utterance2[SEP]utterance3[SEP]"
+    dialogue_len = []  # 记录所有对话tokenize之后的长度，用于统计中位数与均值
+    dialogue_list = []
+    with open(args.save_path, "w", encoding="utf-8") as f:
+        for index, dialogue in enumerate(tqdm(train_data)):
+            if "\r\n" in data:
+                utterances = dialogue.split("\r\n")
+            else:
+                utterances = dialogue.split("\n")
+
+            input_ids = [cls_id]  # 每个dialogue以[CLS]开头
+            for utterance in utterances:
+                input_ids += tokenizer.encode(utterance, add_special_tokens=False)
+                input_ids.append(sep_id)  # 每个utterance之后添加[SEP]，表示utterance结束
+            dialogue_len.append(len(input_ids))
+            dialogue_list.append(input_ids)
+    len_mean = np.mean(dialogue_len)
+    len_median = np.median(dialogue_len)
+    len_max = np.max(dialogue_len)
+    with open(args.save_path, "wb") as f:
+        pickle.dump(dialogue_list, f)
+    logger.info("finish preprocessing data,the result is stored in {}".format(args.save_path))
+    logger.info("mean of dialogue len:{},median of dialogue len:{},max len:{}".format(len_mean, len_median, len_max))
+
+
+if __name__ == '__main__':
+    preprocess()

pytorchtools.py

@@ -0,0 +1,53 @@
+import numpy as np
+import torch
+from os.path import join
+import os
+
+class EarlyStopping:
+    """Early stops the training if validation loss doesn't improve after a given patience."""
+    def __init__(self, patience=7, verbose=False, delta=0, save_path="."):
+        """
+        Args:
+            patience (int): How long to wait after last time validation loss improved.
+                            Default: 7
+            verbose (bool): If True, prints a message for each validation loss improvement.
+                            Default: False
+            delta (float): Minimum change in the monitored quantity to qualify as an improvement.
+                            Default: 0
+        """
+        self.patience = patience
+        self.verbose = verbose
+        self.counter = 0
+        self.best_score = None
+        self.early_stop = False
+        self.val_loss_min = np.Inf
+        self.delta = delta
+        self.save_path = save_path
+
+    def __call__(self, val_loss, model):
+
+        score = -val_loss
+
+        if self.best_score is None:
+            self.best_score = score
+            self.save_checkpoint(val_loss, model)
+        elif score < self.best_score + self.delta:
+            self.counter += 1
+            print(f'EarlyStopping counter: {self.counter} out of {self.patience}')
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.save_checkpoint(val_loss, model)
+            self.counter = 0
+
+    def save_checkpoint(self, val_loss, model):
+        '''Saves model when validation loss decrease.'''
+        if self.verbose:
+            print(f'Validation loss decreased ({self.val_loss_min:.6f} --> {val_loss:.6f}).  Saving model ...')
+        # save_path = join(self.save_path, "best_model")
+        # if not os.path.exists(save_path):
+        #     os.mkdir(save_path)
+        # model_to_save = model.module if hasattr(model, 'module') else model
+        # model_to_save.save_pretrained(save_path)
+        self.val_loss_min = val_loss

requirements.txt

@@ -1,18 +1,2 @@
-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
+torch
+transformers

templates/chat_template.html

@@ -0,0 +1,240 @@
+<html lang="zh">
+    <head>
+        <meta http-equiv="Content-Type" content="text/html; charset=UTF-8">
+        <title>聊天机器人</title>
+        <style>
+            body {
+                padding:0;
+                margin:0;
+                background:-moz-linear-gradient(-45deg,#183850 0,#183850 25%,#192C46 50%,#22254C 75%,#22254C 100%);
+                background:-webkit-linear-gradient(-45deg,#183850 0,#183850 25%,#192C46 50%,#22254C 75%,#22254C 100%);
+                background-repeat:no-repeat;
+                background-attachment:fixed
+            }
+            ::-webkit-scrollbar {
+                width:10px
+            }
+            ::-webkit-scrollbar-track {
+                border-radius:10px;
+                background-color:rgba(25,147,147,0.1)
+            }
+            ::-webkit-scrollbar-thumb {
+                border-radius:10px;
+                background-color:rgba(25,147,147,0.2)
+            }
+            .chat-thread {
+                margin:24px auto 0 auto;
+                padding:0 20px 0 0;
+                list-style:none;
+                overflow-y:scroll;
+                overflow-x:hidden
+            }
+            .chat-thread li {
+                position:relative;
+                clear:both;
+                display:inline-block;
+                padding:16px 40px 16px 20px;
+                margin:0 0 20px 0;
+                font:16px/20px "Noto Sans",sans-serif;
+                border-radius:10px;
+                background-color:rgba(25,147,147,0.2)
+            }
+            .chat-thread li:before {
+                position:absolute;
+                top:0;
+                width:50px;
+                height:50px;
+                border-radius:50px;
+                content:""
+            }
+            .chat-thread li:after {
+                position:absolute;
+                top:15px;
+                content:"";
+                width:0;
+                height:0;
+                border-top:15px solid rgba(25,147,147,0.2)
+            }
+            .chat-thread li:nth-child(odd) {
+                animation:show-chat-odd .15s 1 ease-in;
+                -moz-animation:show-chat-odd .15s 1 ease-in;
+                -webkit-animation:show-chat-odd .15s 1 ease-in;
+                float:right;
+                margin-right:80px;
+                color:#0AD5C1
+            }
+            .chat-thread li:nth-child(odd):before {
+                right:-80px;
+                background-image:url(data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD/4QAiRXhpZgAATU0AKgAAAAgAAQESAAMAAAABAAEAAAAAAAD/2wBDAAIBAQIBAQICAgICAgICAwUDAwMDAwYEBAMFBwYHBwcGBwcICQsJCAgKCAcHCg0KCgsMDAwMBwkODw0MDgsMDAz/2wBDAQICAgMDAwYDAwYMCAcIDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAz/wAARCAAwADADASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD9WfjH8RY/gz8G/FXi6SxuNUj8K6Pd6u9pbusctyIIWlKBm+VcheWbhRkngV/Ip+1z+1Rrv7ZH7RHij4i+I5IbO98Uag96LeKRpIrRGPyRqxwz7F2qGI6IAAAAo/rc/aa8ITfED9mH4keH7eSSG413wnq2nQyIwVo3mspo1YEggEFgQSCBX8o//BPf9k9P2s/jCuj3kzW1jHD58pUcsuR/PNVicRCjTdWeyNcJh516qpQ3Z4/B4wOht5Kw2l1GpKkqpQsPUHP69ah8Q6omtx281v50jIpMxcfNx/e+nY85Hev248D/APBAf4QeJNLjElncblKecwlkEijPODuwc+4r2+5/4IhfAPwZ4Cls7HwkvmX0IjaWR2eSNQQSQzE9eM4xXz8uJqDhzwi393+Z9J/qzXUuSc1+P+R+RX/BGX9rXVv2VP26fhnq0OuSafoOvapB4b8SLI5e3nsLuZY2Eirk/u3MUqnGQyDoM5/qKv7fAYHduXIr+XH9sH9jS7+CH/BQ7Q/h38OzfC417VtMtdFB/eNHPPcIkZTPQiQqeoHBr+pbWGzcTN+7ZixyyDap57DnA9BX0GDxEa1JVobSVz5zG0XRqulPeOhx/wC0s3iOf9nHxwPB+p3Gj+Ko9EuptKvYDtlguEjLoVPZiVwCMYJ6jqPwV/Yw+F+rfCU/FvxZ4V023/tJdaaCxN5ZStDHbeTHct+5g+fav2jACf3V9MV/QxC2VwQrA9mGQR6EV+fPxM/Z/tv2GPja81q1rd+GPG2pXeo2luEZTZoRCrQtuJyU+QAg4IxwOg8XPqdRUnUjqrWa6bpr9Ue/w5Uoyl7Cekua8X11TTV/ua+Z5b+zN/wUq+JkHiqz0jXvhTZtpNxcJpsmpWcWoW/7w+V8wjvIFBA82M5SQn5wBnBx137Sf/BQ/wCNHhj4qt4L8O/De1h0vfLBLq95pV3qUilVclhHCEjjXCOQ0suMqRjOAfV9YbwX4An8OXka+H9Bs9V1CDz7maSO2RlEilYkLY+ZmIIUdcH059Q8W6x4H+Ivi7XNWsX8PeIotNvZBBeWzR3Qhl3F3hYjO11ypK5yAwPevj1Wgpe0Ufd2tf8Ar8j7SWGly+zbbff57f079fI+J/AnwN8SePf+Ci37OPjfVGSx8RW9hqF1PNFYvYCVYrixXa0UpMke+0vL0AH5g+CAo6fqzqH+rYjjivF/2Z/Ctp8SfENx48uods2l3E2l6dFsG1V2Rs8mevJOMdMoD1Ax7JqDARt9K+34fo1I4ZSn12Xlv+rPgOIq1OWI5KfTd+e2/XRL8Sp4l8U6X4E8O3Wsa5qem6Lo+nrvur/ULqO1tbZfWSWQhEHuxFflH/wVU/4K7eA/iR+0X8Nfh78O9S0HxpodrbX1/rHiDT5TcRxXLKRHaQSLhGwtu8khG5W3whWBRgfyz/am/a8+JH7VOrW99468b+KvFqwP5kS6tfl4LUn/AJ4W64ggz38tQT615J4b1a+k+L1vdW5bzNLX7TFAnGUSMb1x/wBcw/4V7uZYJLCzUne6PHy3ESjioSj0aP2U1n40eIPiT4A023tdH8H+ItBkKtJb65bXlxGMfxbbYFmBGRxyPQ549x8CfF/XvDXwW8zWLXwlovh+0jCWVnpGnXlrsXAUbjckFvTPloWJHAr86vgn+1F48+A0EM3hO4WazviJI7e4G9Iy2CWUggr26HFWP25f2/8Axlr/AIANn4g1yO81O8zst7VRHHBJtIUYHzFgCSSTwucckZ/L44KpUaoU0tX8z9YlmFKnSdSa21uffv8AwR3/AOCr+l/Fzx18Ufhz401rQ9Hh07xNczeDLq4uEtYLmzUxW72gd8K0hlQzLli0hnkA4VVr9FtSDLvVlZWUcgjBFfyGaTdtb+DZLKQS3PmKwcSlW8yR+OmSPvkH2r7S/Ym/4LS/Fz9jfwvZeHG1mDxh4VsU2QaZr4e4S2jHGyGVSJYQvZFbyx12Hqf1qhlvJSjGD2SX3H4/iMX7SrKo18Tb+8//2Q==)
+            }
+            .chat-thread li:nth-child(odd):after {
+                border-right:15px solid transparent;
+                right:-15px
+            }
+            .chat-thread li:nth-child(even) {
+                animation:show-chat-even .15s 1 ease-in;
+                -moz-animation:show-chat-even .15s 1 ease-in;
+                -webkit-animation:show-chat-even .15s 1 ease-in;
+                float:left;
+                margin-left:80px;
+                color:#0EC879
+            }
+            .chat-thread li:nth-child(even):before {
+                left:-80px;
+                background-image:url(data:image/jpeg;base64,/9j/4AAQSkZJRgABAQEAYABgAAD/4QAiRXhpZgAATU0AKgAAAAgAAQESAAMAAAABAAEAAAAAAAD/2wBDAAIBAQIBAQICAgICAgICAwUDAwMDAwYEBAMFBwYHBwcGBwcICQsJCAgKCAcHCg0KCgsMDAwMBwkODw0MDgsMDAz/2wBDAQICAgMDAwYDAwYMCAcIDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAwMDAz/wAARCAAwADADASIAAhEBAxEB/8QAHwAAAQUBAQEBAQEAAAAAAAAAAAECAwQFBgcICQoL/8QAtRAAAgEDAwIEAwUFBAQAAAF9AQIDAAQRBRIhMUEGE1FhByJxFDKBkaEII0KxwRVS0fAkM2JyggkKFhcYGRolJicoKSo0NTY3ODk6Q0RFRkdISUpTVFVWV1hZWmNkZWZnaGlqc3R1dnd4eXqDhIWGh4iJipKTlJWWl5iZmqKjpKWmp6ipqrKztLW2t7i5usLDxMXGx8jJytLT1NXW19jZ2uHi4+Tl5ufo6erx8vP09fb3+Pn6/8QAHwEAAwEBAQEBAQEBAQAAAAAAAAECAwQFBgcICQoL/8QAtREAAgECBAQDBAcFBAQAAQJ3AAECAxEEBSExBhJBUQdhcRMiMoEIFEKRobHBCSMzUvAVYnLRChYkNOEl8RcYGRomJygpKjU2Nzg5OkNERUZHSElKU1RVVldYWVpjZGVmZ2hpanN0dXZ3eHl6goOEhYaHiImKkpOUlZaXmJmaoqOkpaanqKmqsrO0tba3uLm6wsPExcbHyMnK0tPU1dbX2Nna4uPk5ebn6Onq8vP09fb3+Pn6/9oADAMBAAIRAxEAPwD9aDlnwuSW6e1fiN/wcL/8FGNW+K/xsm+A/g/VZ7Pwn4PZR4ke3nKjWNQYbvKlI5MUK4+T+JzzwBX7SeOfGtv8M/AeveJrxo47Xw7plxqUrynCqIonk59iVA/Gv5Pzreq/G/4hXWrzLLda1401WXUJVLFnllupi4BPXChgPotZ4mdloehltHmlc0L/AMM3114JMOhxXVxbuB9peBA0l4wwcKoAYbSAAc7eRwvU7PxZ+FF/8Bf2g9f0f7RdWY0vU5re3ubiEtDKiNtDFSo+RuOVxxt+tfuF/wAE2f2IvDvwM+HmnRz6XY3msSIkl1czwK5LEdBuGML/APrr1z9r/wD4J1+Cf2svhRJZ6jZ20OqeU5tL2OELJAT0yQMkZPr714/1+cnotEfRSwVOLSb1fkfmT/wTA/4KL6t+xf8AE+w0zVtUvLj4Z+IriODXNIllaaHRZXIVb62VifLwSC+0gMhyRkZr9wt6OiSRyLNHIodJE+66sMhlPoRyMckdTwK/md+Ofwy179nv4n3XgvxJD5epaPKbMsVwLuBsiNwf4g3Sv2N/4Iaftmt+07+yzN4V1a6a48UfC+VNMnaRsvdWDA/ZpSepK4MZHUbQe4r0MHir+6eTmeDSXPH5/wCZof8ABe74sXvwz/4JgeP4dOby5vEyx6NNKr7TFbu4FwVH8R2Ky49OTivxg/4J4/BtdW8V3njfVpNasNB8HkSzXOkWouLyJ8Ar5SEH1Hbp1IFfrl/wcf6ZPrf/AATU1ZYGk8631SK7kZWO7yY1fzFb+8pTgk8nPtXwb/wRF/aA0X/hO7/w3NHEl5d20U95bOg8ssn7vzVGMYYYB59awzWUuT3ex0cPxg5LmP0A/YX/AGjo/Et60tr4h+JWp6d9sbSZ9L8aaZaW9w0yvEjzWlzbkpcRBpoUJBwGlVQSQceo/tfftRX/AOz98WtL0SXUPiRNDql6limm+D9Ctbnyy+8LLcTXBUQxsysN2dq7QGK7lLX/AIreMtI8J6j8PrWWS3tdP1TW7eWWST5YkEBUpk5xkeYyqAAdrNjjOfoZ9Q0vxFq9xNHHDcT6fdT2ouFwWUhyJEDDH8QIK+g75ryaWt0vI9rEbpvzPx//AOC9PwM/4Sn4I+DfjBp7a5JeRzLZzf2ppwsdQMMnzRmWIKCHVgRyDwcjjBPk/wDwQD+KzeD/APgoEtrAzJZ+P9Ams7iHcQi3EbBs49AwZuTkfKeARX3t/wAF8vij4V8JfsjGz8Sala2Md9dl7K1aJJLjVZ0UlYYwzDHJBLD7oFflj/wQxW/vP+Cl3wzt7WQx26T6hdSxDG3yTp03mbcYGSfLz3yvpXdhLp3OHGOM6fqmfuJ/wUA/Z/j/AGnP2SvHXhXy2kurrRrs2xHVX+zuAR745Hvkda/mj/Zf+OU/7LPx30PxZNBJPZ2f+i6rAv3pIX4kx6srZYfSv6lPjvDfXfw9utLsZJrdtXcWU8yY3RQyHbIRn+EpkZ49ODX8vvxE+CGteO/F3xRt/Cfh+61TTfh9eXlzq0thG08djZJcPGrFlG0IgwuSynI4zXr4qKl7rPDy/mS5o9Hc/aVfjL4u+NvhbwBqXw1l0PWtLup0uLt72zF4UUYMbLGeCpG9WHUEj3r61+A3iXx3aeDobjx1Joq7YlKfZrVrZ4lUAfvixIYgADd1OM9a/Mz/AIJRL4s8K+CvCb6HNNHJdW+yWCVN8TFF3BuowSpPQjoO/Nfon8bNfutL+HdxfeItWiVbG0aZwxEMAYLnkD+FevJ/GvjYS5JSR+hVpqeHjScVfe/XU/Gv/guj+3bZ/tp/tJr4f0GCT/hG/he8ul2M0h2i9vJCPtE/+6OFXPoTXgv/AATw/aGX9lL9sb4d+NpvMls9B1X7PqkSnEj2kwNvcbR/eEMjuP8AajUd66z9qH9jLxx4c8Bx/GCTRb+68O/EK7vNR0kWVu8lxDaRSH/TJ4wCVhkOWR8bSoznBFV/2OP+CaHjf/go1Lcat4DktdF0/Q7iO31vUdWtJ1023YjcRbSxqTPcBcFrcYYblLPGCCfo8PJygmkfJ1qfJN32P//Z)
+            }
+            .chat-thread li:nth-child(even):after {
+                border-left:15px solid transparent;
+                left:-15px
+            }
+            .chat-window {
+                position:fixed;
+                bottom:18px
+            }
+            .chat-window-message {
+                width:100%;
+                height:48px;
+                font:32px/48px "Noto Sans",sans-serif;
+                background:0;
+                color:#0AD5C1;
+                border:0;
+                border-bottom:1px solid rgba(25,147,147,0.2);
+                outline:0
+            }
+            @media all and (max-width:767px) {
+                .chat-thread {
+                    width:90%;
+                    height:90%
+                }
+                .chat-window {
+                    left:5%;
+                    width:90%
+                }
+            }
+            @media all and (min-width:768px) {
+                .chat-thread {
+                    width:50%;
+                    height:90%
+                }
+                .chat-window {
+                    left:25%;
+                    width:50%
+                }
+            }
+            @keyframes show-chat-even {
+                0% {
+                    margin-left:-480px
+                }
+                100% {
+                    margin-left:0
+                }
+            }
+            @-moz-keyframes show-chat-even {
+                0% {
+                    margin-left:-480px
+                }
+                100% {
+                    margin-left:0
+                }
+            }
+            @-webkit-keyframes show-chat-even {
+                0% {
+                    margin-left:-480px
+                }
+                100% {
+                    margin-left:0
+                }
+            }
+            @keyframes show-chat-odd {
+                0% {
+                    margin-right:-480px
+                }
+                100% {
+                    margin-right:0
+                }
+            }
+            @-moz-keyframes show-chat-odd {
+                0% {
+                    margin-right:-480px
+                }
+                100% {
+                    margin-right:0
+                }
+            }
+            @-webkit-keyframes show-chat-odd {
+                0% {
+                    margin-right:-480px
+                }
+                100% {
+                    margin-right:0
+                }
+            }
+        </style>
+    </head>
+    <body onload="loadhistory()">
+        <ul class="chat-thread">
+
+        </ul>
+        <div class="chat-window">
+            <input class="chat-window-message" name="chat-window-message" type="text" autocomplete="off" autofocus="" placeholder="对我说HELP，看看我能干什么~">
+        </div>
+        <script src="https://cdn.bootcdn.net/ajax/libs/jquery/2.1.3/jquery.min.js"></script>
+        <script>
+            var chat_window = document.querySelector(".chat-window");
+
+            chat_window.onkeydown=function(event){
+                var e = event || window.event || arguments.callee.caller.arguments[0];
+                if (e && e.keyCode == 13 ) {
+                    send_data();
+                }
+            }
+
+            function send_data() {
+                var chat_thread = document.querySelector(".chat-thread");
+                var chat_window_message = document.querySelector(".chat-window-message");
+                chat_window_message.disabled = true;
+                var text = chat_window_message.value;
+                var new_li_label = document.createElement("li"), new_li_text = document.createTextNode(text);
+            	new_li_label.appendChild(new_li_text);
+            	chat_thread.appendChild(new_li_label);
+            	chat_thread.scrollTop = chat_thread.scrollHeight;
+                chat_window_message.value = "";
+                document.title = "聊天机器人 ~ 对方正在输入…"
+                $.getJSON("/chitchat/chat?text="+text, function(data){
+                    var new_li_label = document.createElement("li");
+                    data.forEach(function(item){
+                        var new_text = document.createTextNode(item);
+                        var new_span = document.createElement("span");
+                        new_span.appendChild(new_text);
+                        var new_br = document.createElement("br");
+                        new_li_label.appendChild(new_span);
+                        new_li_label.appendChild(new_br);
+                    });
+                    chat_thread.appendChild(new_li_label);
+                    chat_thread.scrollTop = chat_thread.scrollHeight;
+
+                    document.title = "聊天机器人"
+                    chat_window_message.disabled = false;
+                });
+            }
+
+            function loadhistory() {
+                var chat_thread = document.querySelector(".chat-thread");
+                var chat_window_message = document.querySelector(".chat-window-message");
+                chat_window_message.disabled = true;
+                document.title = "聊天机器人 ~ 正在回忆…"
+                $.getJSON("/chitchat/history", function(data){
+                    data.forEach(function(item) {
+                            var new_li_label = document.createElement("li"),new_li_text = document.createTextNode(item);
+                            new_li_label.appendChild(new_li_text);
+                            chat_thread.appendChild(new_li_label);
+                        });
+                        chat_thread.scrollTop = chat_thread.scrollHeight;
+                        
+                        chat_window_message.disabled = false;
+                        document.title = "聊天机器人"
+                });
+            }
+        </script>
+    </body>
+</html>

train.py

@@ -0,0 +1,432 @@
+import argparse
+import math
+import time
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+import logging
+from datetime import datetime
+import os
+from torch.utils.data import Dataset, DataLoader
+from os.path import join, exists
+from torch.nn import CrossEntropyLoss
+from tqdm import tqdm
+from torch.nn import DataParallel
+import transformers
+import pickle
+import sys
+from pytorchtools import EarlyStopping
+from sklearn.model_selection import train_test_split
+from data_parallel import BalancedDataParallel
+from transformers import GPT2TokenizerFast, GPT2LMHeadModel, GPT2Config
+from transformers import BertTokenizerFast
+import pandas as pd
+import torch.nn.utils.rnn as rnn_utils
+import numpy as np
+from dataset import MyDataset
+from chatbot.config import config
+
+
+def set_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--device', default='3', type=str, required=False, help='设置使用哪些显卡')
+    parser.add_argument('--no_cuda', action='store_true', help='不使用GPU进行训练')
+    parser.add_argument('--vocab_path', default='vocab/vocab.txt', type=str, required=False,
+                        help='词表路径')
+    parser.add_argument('--model_config', default='config/config.json', type=str, required=False,
+                        help='设置模型参数')
+    parser.add_argument('--train_path', default='data/train.pkl', type=str, required=False, help='训练集路径')
+    parser.add_argument('--max_len', default=150, type=int, required=False, help='训练时，输入数据的最大长度')
+
+    parser.add_argument('--log_path', default='data/train.log', type=str, required=False, help='训练日志存放位置')
+    parser.add_argument('--log', default=True, help="是否记录日志")
+    parser.add_argument('--ignore_index', default=-100, type=int, required=False, help='对于ignore_index的label token不计算梯度')
+    # parser.add_argument('--input_len', default=200, type=int, required=False, help='输入的长度')
+    parser.add_argument('--epochs', default=100, type=int, required=False, help='训练的最大轮次')
+    parser.add_argument('--batch_size', default=4, type=int, required=False, help='训练的batch size')
+    parser.add_argument('--gpu0_bsz', default=10, type=int, required=False, help='0号卡的batch size')
+    parser.add_argument('--lr', default=2.6e-5, type=float, required=False, help='学习率')
+    parser.add_argument('--eps', default=1.0e-09, type=float, required=False, help='衰减率')
+    parser.add_argument('--log_step', default=1, type=int, required=False, help='多少步汇报一次loss')
+    parser.add_argument('--gradient_accumulation_steps', default=4, type=int, required=False, help='梯度积累')
+    parser.add_argument('--max_grad_norm', default=2.0, type=float, required=False)
+    parser.add_argument('--save_model_path', default='model', type=str, required=False,
+                        help='模型输出路径')
+    parser.add_argument('--pretrained_model', default='', type=str, required=False,
+                        help='预训练的模型的路径')
+    # parser.add_argument('--seed', type=int, default=None, help='设置种子用于生成随机数，以使得训练的结果是确定的')
+    parser.add_argument('--num_workers', type=int, default=0, help="dataloader加载数据时使用的线程数量")
+    parser.add_argument('--patience', type=int, default=0, help="用于early stopping,设为0时,不进行early stopping.early stop得到的模型的生成效果不一定会更好。")
+    parser.add_argument('--warmup_steps', type=int, default=4000, help='warm up步数')
+    # parser.add_argument('--label_smoothing', default=True, action='store_true', help='是否进行标签平滑')
+    parser.add_argument('--val_num', type=int, default=8000, help='验证集大小')
+    args = parser.parse_args()
+    return args
+
+
+def create_logger(args):
+    """
+    将日志输出到日志文件和控制台
+    """
+    logger = logging.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+
+    formatter = logging.Formatter(
+        '%(asctime)s - %(levelname)s - %(message)s')
+
+    # 创建一个handler，用于写入日志文件
+    file_handler = logging.FileHandler(
+        filename=args.log_path)
+    file_handler.setFormatter(formatter)
+    file_handler.setLevel(logging.INFO)
+    logger.addHandler(file_handler)
+
+    # 创建一个handler，用于将日志输出到控制台
+    console = logging.StreamHandler()
+    console.setLevel(logging.DEBUG)
+    console.setFormatter(formatter)
+    logger.addHandler(console)
+
+    return logger
+
+
+def collate_fn(batch):
+    input_ids = rnn_utils.pad_sequence(batch, batch_first=True, padding_value=0)
+    labels = rnn_utils.pad_sequence(batch, batch_first=True, padding_value=-100)
+    return input_ids, labels
+
+
+# def padding_batch(data_list, pad_id):
+#     """
+#     使用pad_id将data_list的每条数据，填充至data_list中最长的长度
+#     :param data_list:
+#     :param pad_id:
+#     :return:
+#     """
+#     # 统计data_list中的最大长度
+#     max_len = 0
+#     for data in data_list:
+#         max_len = max_len if max_len > len(data) else len(data)
+#
+#     # 对数据进行padding
+#     new_data_list = []
+#     for data in data_list:
+#         new_data = data + [pad_id] * (max_len - len(data))
+#         new_data_list.append(new_data)
+#     return new_data_list
+
+
+def load_dataset(logger, args):
+    """
+    加载训练集和验证集
+    """
+    logger.info("loading training dataset and validating dataset")
+    train_path = args.train_path
+
+    with open(train_path, "rb") as f:
+        input_list = pickle.load(f)
+
+    # 划分训练集与验证集
+    val_num = args.val_num
+    input_list_train = input_list[val_num:]
+    input_list_val = input_list[:val_num]
+    # test
+    # input_list_train = input_list_train[:24]
+    # input_list_val = input_list_val[:24]
+
+    train_dataset = MyDataset(input_list_train, args.max_len)
+    val_dataset = MyDataset(input_list_val, args.max_len)
+
+    return train_dataset, val_dataset
+
+
+def train_epoch(model, train_dataloader, optimizer, scheduler, logger,
+                epoch, args):
+    model.train()
+    device = args.device
+    # pad_id = args.pad_id
+    # sep_id = args.sep_id
+    ignore_index = args.ignore_index
+    epoch_start_time = datetime.now()
+    total_loss = 0  # 记录下整个epoch的loss的总和
+
+    # epoch_correct_num:每个epoch中,output预测正确的word的数量
+    # epoch_total_num: 每个epoch中,output预测的word的总数量
+    epoch_correct_num, epoch_total_num = 0, 0
+
+    for batch_idx, (input_ids, labels) in enumerate(train_dataloader):
+        # 捕获cuda out of memory exception
+        try:
+            input_ids = input_ids.to(device)
+            labels = labels.to(device)
+            outputs = model.forward(input_ids, labels=labels)
+            logits = outputs.logits
+            loss = outputs.loss
+            loss = loss.mean()
+
+            # 统计该batch的预测token的正确数与总数
+            batch_correct_num, batch_total_num = calculate_acc(logits, labels, ignore_index=ignore_index)
+            # 统计该epoch的预测token的正确数与总数
+            epoch_correct_num += batch_correct_num
+            epoch_total_num += batch_total_num
+            # 计算该batch的accuracy
+            batch_acc = batch_correct_num / batch_total_num
+
+            total_loss += loss.item()
+            if args.gradient_accumulation_steps > 1:
+                loss = loss / args.gradient_accumulation_steps
+
+            loss.backward()
+            # 梯度裁剪
+            torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+
+            # 进行一定step的梯度累计之后，更新参数
+            if (batch_idx + 1) % args.gradient_accumulation_steps == 0:
+                # 更新参数
+                optimizer.step()
+                # 更新学习率
+                scheduler.step()
+                # 清空梯度信息
+                optimizer.zero_grad()
+
+            if (batch_idx + 1) % args.log_step == 0:
+                logger.info(
+                    "batch {} of epoch {}, loss {}, batch_acc {}, lr {}".format(
+                        batch_idx + 1, epoch + 1, loss.item() * args.gradient_accumulation_steps, batch_acc, scheduler.get_lr()))
+
+            del input_ids, outputs
+
+        except RuntimeError as exception:
+            if "out of memory" in str(exception):
+                logger.info("WARNING: ran out of memory")
+                if hasattr(torch.cuda, 'empty_cache'):
+                    torch.cuda.empty_cache()
+            else:
+                logger.info(str(exception))
+                raise exception
+
+    # 记录当前epoch的平均loss与accuracy
+    epoch_mean_loss = total_loss / len(train_dataloader)
+    epoch_mean_acc = epoch_correct_num / epoch_total_num
+    logger.info(
+        "epoch {}: loss {}, predict_acc {}".format(epoch + 1, epoch_mean_loss, epoch_mean_acc))
+
+    # save model
+    logger.info('saving model for epoch {}'.format(epoch + 1))
+    model_path = join(args.save_model_path, 'epoch{}'.format(epoch + 1))
+    if not os.path.exists(model_path):
+        os.mkdir(model_path)
+    model_to_save = model.module if hasattr(model, 'module') else model
+    model_to_save.save_pretrained(model_path)
+    logger.info('epoch {} finished'.format(epoch + 1))
+    epoch_finish_time = datetime.now()
+    logger.info('time for one epoch: {}'.format(epoch_finish_time - epoch_start_time))
+
+    return epoch_mean_loss
+
+
+def validate_epoch(model, validate_dataloader, logger, epoch, args):
+    logger.info("start validating")
+    model.eval()
+    device = args.device
+    # pad_id = args.pad_id
+    # sep_id = args.sep_id
+    ignore_index = args.ignore_index
+    epoch_start_time = datetime.now()
+    total_loss = 0
+    # 捕获cuda out of memory exception
+    try:
+        with torch.no_grad():
+            for batch_idx, (input_ids, labels) in enumerate(validate_dataloader):
+                input_ids = input_ids.to(device)
+                labels = labels.to(device)
+                outputs = model.forward(input_ids, labels=labels)
+                logits = outputs.logits
+                loss = outputs.loss
+                loss = loss.mean()
+
+                total_loss += loss.item()
+                del input_ids, outputs
+
+            # 记录当前epoch的平均loss
+            epoch_mean_loss = total_loss / len(validate_dataloader)
+            logger.info(
+                "validate epoch {}: loss {}".format(epoch+1, epoch_mean_loss))
+            epoch_finish_time = datetime.now()
+            logger.info('time for validating one epoch: {}'.format(epoch_finish_time - epoch_start_time))
+            return epoch_mean_loss
+    except RuntimeError as exception:
+        if "out of memory" in str(exception):
+            logger.info("WARNING: ran out of memory")
+            if hasattr(torch.cuda, 'empty_cache'):
+                torch.cuda.empty_cache()
+        else:
+            logger.info(str(exception))
+            raise exception
+
+
+def train(model, logger, train_dataset, validate_dataset, args):
+    train_dataloader = DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, collate_fn=collate_fn,
+        drop_last=True
+    )
+    validate_dataloader = DataLoader(validate_dataset, batch_size=args.batch_size, shuffle=True,
+                                     num_workers=args.num_workers, collate_fn=collate_fn, drop_last=True)
+    early_stopping = EarlyStopping(args.patience, verbose=True, save_path=args.save_model_path)
+    t_total = len(train_dataloader) // args.gradient_accumulation_steps * args.epochs
+    optimizer = transformers.AdamW(model.parameters(), lr=args.lr, eps=args.eps)
+    # scheduler = transformers.WarmupLinearSchedule(optimizer, warmup_steps=args.warmup_steps, t_total=t_total)
+    scheduler = transformers.get_linear_schedule_with_warmup(
+        optimizer, num_warmup_steps=args.warmup_steps, num_training_steps=t_total
+    )
+
+    logger.info('starting training')
+
+    # 用于记录每个epoch训练和验证的loss
+    train_losses, validate_losses = [], []
+    # 记录验证集的最小loss
+    best_val_loss = 10000
+    # 开始训练
+    for epoch in range(args.epochs):
+        # ========== train ========== #
+        train_loss = train_epoch(
+            model=model, train_dataloader=train_dataloader,
+            optimizer=optimizer, scheduler=scheduler,
+            logger=logger, epoch=epoch, args=args)
+        train_losses.append(train_loss)
+
+        # ========== validate ========== #
+        validate_loss = validate_epoch(
+            model=model, validate_dataloader=validate_dataloader,
+            logger=logger, epoch=epoch, args=args)
+        validate_losses.append(validate_loss)
+
+        # 保存当前困惑度最低的模型，困惑度低，模型的生成效果不一定会越好
+        if validate_loss < best_val_loss:
+            best_val_loss = validate_loss
+            logger.info('saving current best model for epoch {}'.format(epoch + 1))
+            model_path = join(args.save_model_path, 'min_ppl_model'.format(epoch + 1))
+            if not os.path.exists(model_path):
+                os.mkdir(model_path)
+            model_to_save = model.module if hasattr(model, 'module') else model
+            model_to_save.save_pretrained(model_path)
+
+        #  如果patience=0,则不进行early stopping
+        if args.patience == 0:
+            continue
+        early_stopping(validate_loss, model)
+        if early_stopping.early_stop:
+            logger.info("Early stopping")
+            break
+    logger.info('training finished')
+    logger.info("train_losses:{}".format(train_losses))
+    logger.info("validate_losses:{}".format(validate_losses))
+
+
+def caculate_loss(logit, target, pad_idx, smoothing=True):
+    if smoothing:
+        logit = logit[..., :-1, :].contiguous().view(-1, logit.size(2))
+        target = target[..., 1:].contiguous().view(-1)
+
+        eps = 0.1
+        n_class = logit.size(-1)
+
+        one_hot = torch.zeros_like(logit).scatter(1, target.view(-1, 1), 1)
+        one_hot = one_hot * (1 - eps) + (1 - one_hot) * eps / (n_class - 1)
+        log_prb = F.log_softmax(logit, dim=1)
+
+        non_pad_mask = target.ne(pad_idx)
+        loss = -(one_hot * log_prb).sum(dim=1)
+        loss = loss.masked_select(non_pad_mask).mean()  # average later
+    else:
+        # loss = F.cross_entropy(predict_logit, target, ignore_index=pad_idx)
+        logit = logit[..., :-1, :].contiguous().view(-1, logit.size(-1))
+        labels = target[..., 1:].contiguous().view(-1)
+        loss = F.cross_entropy(logit, labels, ignore_index=pad_idx)
+    return loss
+
+
+def calculate_acc(logit, labels, ignore_index=-100):
+    logit = logit[..., :-1, :].contiguous().view(-1, logit.size(-1))
+    labels = labels[..., 1:].contiguous().view(-1)
+
+    _, logit = logit.max(dim=-1)  # 对于每条数据，返回最大的index
+    # 进行��运算，返回一个tensor，若labels的第i个位置为pad_id，则置为0，否则为1
+    non_pad_mask = labels.ne(ignore_index)
+    n_correct = logit.eq(labels).masked_select(non_pad_mask).sum().item()
+    n_word = non_pad_mask.sum().item()
+    return n_correct, n_word
+
+
+def main():
+    # 初始化参数
+    args = set_args()
+
+    # 设置使用哪些显卡进行训练
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
+
+    args.cuda = not args.no_cuda
+
+    if args.batch_size < 2048 and args.warmup_steps <= 4000:
+        print('[Warning] The warmup steps may be not enough.\n' \
+              '(sz_b, warmup) = (2048, 4000) is the official setting.\n' \
+              'Using smaller batch w/o longer warmup may cause ' \
+              'the warmup stage ends with only little data trained.')
+
+    # 创建日志对象
+    logger = create_logger(args)
+    # 当用户使用GPU,并且GPU可用时
+    args.cuda = torch.cuda.is_available() and not args.no_cuda
+    device = 'cuda:0' if args.cuda else 'cpu'
+    args.device = device
+    logger.info('using device:{}'.format(device))
+
+    # 初始化tokenizer
+    tokenizer = BertTokenizerFast(vocab_file=args.vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
+    special_tokens = []
+    for key in config["mask_token"].keys():
+      special_tokens.append(key)
+    tokenizer.add_special_tokens( {'additional_special_tokens':special_tokens} )
+    args.sep_id = tokenizer.sep_token_id
+    args.pad_id = tokenizer.pad_token_id
+    args.cls_id = tokenizer.cls_token_id
+
+    # 创建模型的输出目录
+    if not os.path.exists(args.save_model_path):
+        os.mkdir(args.save_model_path)
+
+    # 创建模型
+    if args.pretrained_model:  # 加载预训练模型
+        model = GPT2LMHeadModel.from_pretrained(args.pretrained_model)
+    else:  # 初始化模型
+        model_config = GPT2Config.from_json_file(args.model_config)
+        model = GPT2LMHeadModel(config=model_config)
+    model = model.to(device)
+    logger.info('model config:\n{}'.format(model.config.to_json_string()))
+    assert model.config.vocab_size == tokenizer.vocab_size
+
+    # 并行训练模型
+    if args.cuda and torch.cuda.device_count() > 1:
+        model = DataParallel(model).cuda()
+        # model = BalancedDataParallel(args.gpu0_bsz, model, dim=0).cuda()
+        logger.info("use GPU {} to train".format(args.device))
+
+    # 计算模型参数数量
+    num_parameters = 0
+    parameters = model.parameters()
+    for parameter in parameters:
+        num_parameters += parameter.numel()
+    logger.info('number of model parameters: {}'.format(num_parameters))
+
+    # 记录参数设置
+    logger.info("args:{}".format(args))
+
+    # 加载训练集和验证集
+    # ========= Loading Dataset ========= #
+    train_dataset, validate_dataset = load_dataset(logger, args)
+
+    train(model, logger, train_dataset, validate_dataset, args)
+
+
+if __name__ == '__main__':
+    main()

web.py

@@ -0,0 +1,151 @@
+
+import os
+import random
+import re
+import requests
+import argparse
+import string
+
+from datetime import timedelta
+from flask import Flask, session, request, jsonify, render_template
+
+from transformers.models.bert.tokenization_bert import BertTokenizer
+
+from bot.chatbot import ChatBot
+from bot.config import special_token_list
+
+app = Flask(__name__)
+app.config["SECRET_KEY"] = os.urandom(74)
+app.config["PERMANENT_SESSION_LIFETIME"] = timedelta(days=7)
+
+tokenizer:BertTokenizer = None
+
+history_matrix:dict = {}
+
+def move_history_from_session_to_global_memory() -> None:
+    global history_matrix
+
+    if session.get( "session_hash") and session["history"]:
+        history_matrix[session["session_hash"]] = session["history"]
+
+def move_history_from_global_memory_to_session() -> None:
+    global history_matrix
+
+    if session.get( "session_hash"):
+        session["history"] = history_matrix.get( session.get( "session_hash") )
+
+def set_args() -> argparse.Namespace:
+    parser:argparse.ArgumentParser = argparse.ArgumentParser()
+    parser.add_argument("--vocab_path", default=None, type=str, required=False, help="选择词库")
+    parser.add_argument("--model_path", default="lewiswu1209/Winnie", type=str, required=False, help="对话模型路径")
+
+    return parser.parse_args()
+
+@app.route("/chitchat/history", methods = ["GET"])
+def get_history_list() -> str:
+    global tokenizer
+
+    move_history_from_global_memory_to_session()
+
+    history_list:list = session.get("history")
+    if history_list is None:
+        history_list = []
+
+    history:list = []
+    for history_ids in history_list:
+        tokens = tokenizer.convert_ids_to_tokens(history_ids)
+        fixed_tokens = []
+        for token in tokens:
+            if token.startswith("##"):
+                token = token[2:]
+            fixed_tokens.append(token)
+        history.append( "".join( fixed_tokens ) )
+
+    return jsonify(history)
+
+@app.route("/chitchat/chat", methods = ["GET"])
+def talk() -> str:
+    global tokenizer
+    global history_matrix
+
+    if request.args.get("hash"):
+        session["session_hash"] = request.args.get("hash")
+        move_history_from_global_memory_to_session()
+
+    if session.get("session_hash") is None:
+        session["session_hash"] = "".join( random.sample(string.ascii_lowercase + string.digits, 11) )
+
+    if request.args.get("text"):
+        input_text = request.args.get("text")
+        history_list = session.get("history")
+
+        if input_text.upper()=="HELP":
+            help_info_list = ["输入任意文字，Winnie会回答你的问题",
+            "输入ERASE MEMORY，Winnie会清空记忆",
+            "输入\"<TAG>=<VALUE>\"，Winnie会记录你的角色信息",
+            "例如：<NAME>=Vicky，Winnie会修改自己的名字",
+            "可以修改的角色信息有：",
+            "<NAME>, <GENDER>, <YEAROFBIRTH>, <MONTHOFBIRTH>, <DAYOFBIRTH>, <ZODIAC>, <AGE>",
+            "输入“上联：XXXXXXX”，Winnie会和你对对联",
+            "输入“写诗：XXXXXXX”，Winnie会以XXXXXXX为开头写诗"
+            ]
+            return jsonify(help_info_list)
+
+        if history_list is None or len(history_list)==0 or input_text == "ERASE MEMORY":
+            history_list = []
+            output_text = requests.post(
+                url='https://hf.space/embed/lewiswu1209/Winnie/+/api/predict/',
+                json={"data": ["ERASE MEMORY"], "session_hash": session["session_hash"]}
+            ).json()["data"][0]
+
+        if input_text != "ERASE MEMORY":
+            if not re.match( r"^<.+>=.+$", input_text ):
+                history_list.append( tokenizer.encode(input_text, add_special_tokens=False) )
+            output_text = requests.post(
+                url='https://hf.space/embed/lewiswu1209/Winnie/+/api/predict/',
+                json={"data": [input_text], "session_hash": session["session_hash"]}
+            ).json()["data"][0]
+            if not re.match( r"^<.+>=.+$", input_text ):
+                history_list.append( tokenizer.encode(output_text, add_special_tokens=False) )
+
+        session["history"] = history_list
+        history_matrix[session["session_hash"]] = history_list
+        return jsonify([output_text])
+    else:
+        return jsonify([""])
+
+@app.route("/")
+def index() -> str:
+  return "Hello world!"
+
+@app.route("/chitchat/hash", methods = ["GET"])
+def get_hash() -> str:
+    global history_matrix
+
+    if request.args.get("hash"):
+        session["session_hash"] = request.args.get("hash")
+        move_history_from_global_memory_to_session()
+    hash = session.get("session_hash")
+    if hash:
+        return session.get("session_hash")
+    else:
+        return " "
+
+@app.route( "/chitchat", methods = ["GET"] )
+def chitchat() -> str:
+    return render_template( "chat_template.html" )
+
+def main() -> None:
+    global tokenizer
+
+    args = set_args()
+    tokenizer = ChatBot.get_tokenizer(
+        args.model_path,
+        vocab_path=args.vocab_path,
+        special_token_list = special_token_list
+    )
+
+    app.run( host = "127.0.0.1", port = 8080 )
+
+if __name__ == "__main__":
+    main()