Merge branch 'shubaoyu-tombyshen' into main

add suport for multi gpu and plain beam search

Author: wxl1999

crslab/config/config.py

@@ -36,6 +36,10 @@ def __init__(self, config_file, gpu='-1', debug=False):
         self.opt = self.load_yaml_configs(config_file)
         # gpu
         os.environ['CUDA_VISIBLE_DEVICES'] = gpu
+        gpu = gpu.split(",")
+        for i in range(len(gpu)):
+            gpu[i] = int(gpu[i])
+        self.opt["gpu"] = gpu
         # dataset
         dataset = self.opt['dataset']
         tokenize = self.opt['tokenize']

crslab/model/__init__.py

@@ -8,6 +8,7 @@
 # @Email  : francis_kun_zhou@163.com, wxl1999@foxmail.com
 
 from loguru import logger
+import torch
 
 from .conversation import *
 from .crs import *
@@ -43,6 +44,13 @@ def get_model(config, model_name, device, vocab, side_data=None):
     if model_name in Model_register_table:
         model = Model_register_table[model_name](config, device, vocab, side_data)
         logger.info(f'[Build model {model_name}]')
-        return model
+        if config.opt["gpu"] == [-1]:
+            return model
+        else:
+            if len(config.opt["gpu"]) > 1 and model_name == 'PMI':
+                logger.info(f'[PMI model does not support multi GPUs yet, using single GPU now]')
+                return model.to(device)
+            return torch.nn.DataParallel(model, device_ids=config["gpu"])
+
     else:
         raise NotImplementedError('Model [{}] has not been implemented'.format(model_name))

crslab/model/conversation/gpt2/gpt2.py

@@ -64,9 +64,8 @@ def build_model(self):
         self.model = GPT2LMHeadModel.from_pretrained(self.dpath)
         self.loss = CrossEntropyLoss(ignore_index=self.pad_id)
 
-    def converse(self, batch, mode):
+    def forward(self, batch, mode):
         _, _, input_ids, context, _, _, y = batch
-
         if mode != 'test':
             # torch.tensor's shape = (bs, seq_len, v_s); tuple's length = 12
             lm_logits = self.model(input_ids).logits
@@ -119,3 +118,44 @@ def calculate_loss(self, logit, labels):
 
         loss = self.loss(logit.reshape(-1, logit.size(-1)), labels.reshape(-1))
         return loss
+
+    def generate_bs(self, context, beam=4):
+        context = context[..., -self.response_truncate + 1:]
+        context_former = context
+        batch_size = context.shape[0]
+        sequences = [[[list(), 1.0]]] * batch_size
+        for i in range(self.response_truncate - 1):
+            if sequences != [[[list(), 1.0]]] * batch_size:
+                context = []
+                for i in range(batch_size):
+                    for cand in sequences[i]:
+                        text = torch.cat((context_former[i], torch.tensor(cand[0]).to(self.device)))  # 由于取消了state，与之前的context拼接
+                        context.append(text)
+                context = torch.stack(context)
+            with torch.no_grad():
+                outputs = self.model(context)
+            last_hidden_state, state = outputs.logits, outputs.past_key_values
+            next_token_logits = last_hidden_state[:, -1, :]
+            next_token_probs = torch.nn.functional.softmax(next_token_logits)
+            topk = torch.topk(next_token_probs, beam, dim=-1)
+            probs = topk.values.reshape([batch_size, -1, beam])  # (bs, candidate, beam)
+            preds = topk.indices.reshape([batch_size, -1, beam])  # (bs, candidate, beam)
+
+            for j in range(batch_size):
+                all_candidates = []
+                for n in range(len(sequences[j])):
+                    for k in range(beam):
+                        seq = sequences[j][n][0]
+                        prob = sequences[j][n][1]
+                        seq_tmp = seq.copy()
+                        seq_tmp.append(preds[j][n][k])
+                        candidate = [seq_tmp, prob * probs[j][n][k]]
+                        all_candidates.append(candidate)
+                ordered = sorted(all_candidates, key=lambda tup: tup[1], reverse=True)
+                sequences[j] = ordered[:beam]
+
+        res = []
+        for i in range(batch_size):
+            res.append(torch.stack(sequences[i][0][0]))
+        res = torch.stack(res)
+        return res

crslab/model/conversation/transformer/transformer.py

@@ -190,7 +190,61 @@ def _decode_greedy_with_kg(self, token_encoding):
         logits = torch.cat(logits, dim=1)
         return logits, inputs
 
-    def converse(self, batch, mode):
+    def _decode_beam_search_with_kg(self, token_encoding, beam=4):
+        batch_size = token_encoding[0].shape[0]
+        xs = self._starts(batch_size).long().reshape(1, batch_size, -1)
+        incr_state = None
+        sequences = [[[list(), list(), 1.0]]] * batch_size
+        for i in range(self.longest_label):
+            # at beginning there is 1 candidate, when i!=0 there are 4 candidates
+            if i == 1:
+                token_encoding = (token_encoding[0].repeat(beam, 1, 1),
+                                  token_encoding[1].repeat(beam, 1, 1))
+            if i != 0:
+                xs = []
+                for d in range(len(sequences[0])):
+                    for j in range(batch_size):
+                        text = sequences[j][d][0]
+                        xs.append(text)
+                xs = torch.stack(xs).reshape(beam, batch_size, -1)  # (beam, batch_size, _)
+
+            dialog_latent, incr_state = self.conv_decoder(xs.reshape(len(sequences[0]) * batch_size, -1),
+                                                          token_encoding,
+                                                          incr_state)
+            dialog_latent = dialog_latent[:, -1:, :]  # (bs, 1, dim)
+            gen_logits = F.linear(dialog_latent, self.token_embedding.weight)
+
+            logits = gen_logits.reshape(len(sequences[0]), batch_size, 1, -1)
+            # turn into probabilities,in case of negative numbers
+            probs, preds = torch.nn.functional.softmax(logits).topk(beam, dim=-1)
+
+            # (candeidate, bs, 1 , beam) during first loop, candidate=1, otherwise candidate=beam
+
+            for j in range(batch_size):
+                all_candidates = []
+                for n in range(len(sequences[j])):
+                    for k in range(beam):
+                        prob = sequences[j][n][2]
+                        logit = sequences[j][n][1]
+                        if logit == []:
+                            logit_tmp = logits[n][j][0].unsqueeze(0)
+                        else:
+                            logit_tmp = torch.cat((logit, logits[n][j][0].unsqueeze(0)), dim=0)
+                        seq_tmp = torch.cat((xs[n][j].reshape(-1), preds[n][j][0][k].reshape(-1)))
+                        candidate = [seq_tmp, logit_tmp, prob * probs[n][j][0][k]]
+                        all_candidates.append(candidate)
+                ordered = sorted(all_candidates, key=lambda tup: tup[2], reverse=True)
+                sequences[j] = ordered[:beam]
+
+            # check if everyone has generated an end token
+            all_finished = ((xs == self.end_token_idx).sum(dim=1) > 0).sum().item() == batch_size
+            if all_finished:
+                break
+        logits = torch.stack([seq[0][1] for seq in sequences])
+        xs = torch.stack([seq[0][0] for seq in sequences])
+        return logits, xs
+
+    def forward(self, batch, mode):
         context_tokens, context_entities, context_words, response = batch
 
         # encoder-decoder

crslab/model/crs/kbrd/kbrd.py

@@ -71,6 +71,8 @@ def __init__(self, opt, device, vocab, side_data):
             side_data (dict): A dictionary record the side data.
 
         """
+        self.device = device
+        self.gpu = opt.get("gpu", -1)
         # vocab
         self.pad_token_idx = vocab['pad']
         self.start_token_idx = vocab['start']
@@ -172,7 +174,7 @@ def _build_conversation_layer(self):
     def encode_user(self, entity_lists, kg_embedding):
         user_repr_list = []
         for entity_list in entity_lists:
-            if not entity_list:
+            if entity_list is not None:
                 user_repr_list.append(torch.zeros(self.user_emb_dim, device=self.device))
                 continue
             user_repr = kg_embedding[entity_list]
@@ -205,17 +207,18 @@ def decode_forced(self, encoder_states, user_embedding, resp):
         return sum_logits, preds
 
     def decode_greedy(self, encoder_states, user_embedding):
+
         bsz = encoder_states[0].shape[0]
         xs = self._starts(bsz)
         incr_state = None
         logits = []
         for i in range(self.longest_label):
-            scores, incr_state = self.decoder(xs, encoder_states, incr_state)
+            scores, incr_state = self.decoder(xs, encoder_states, incr_state)  # incr_state is always None
             scores = scores[:, -1:, :]
             token_logits = F.linear(scores, self.token_embedding.weight)
             user_logits = self.user_proj_2(torch.relu(self.user_proj_1(user_embedding))).unsqueeze(1)
             sum_logits = token_logits + user_logits
-            _, preds = sum_logits.max(dim=-1)
+            probs, preds = sum_logits.max(dim=-1)
             logits.append(scores)
             xs = torch.cat([xs, preds], dim=1)
             # check if everyone has generated an end token
@@ -225,6 +228,62 @@ def decode_greedy(self, encoder_states, user_embedding):
         logits = torch.cat(logits, 1)
         return logits, xs
 
+    def decode_beam_search(self, encoder_states, user_embedding, beam=4):
+        bsz = encoder_states[0].shape[0]
+        xs = self._starts(bsz).reshape(1, bsz, -1)  # (batch_size, _)
+        sequences = [[[list(), list(), 1.0]]] * bsz
+        for i in range(self.longest_label):
+            # at beginning there is 1 candidate, when i!=0 there are 4 candidates
+            if i != 0:
+                xs = []
+                for d in range(len(sequences[0])):
+                    for j in range(bsz):
+                        text = sequences[j][d][0]
+                        xs.append(text)
+                xs = torch.stack(xs).reshape(beam, bsz, -1)  # (beam, batch_size, _)
+
+            with torch.no_grad():
+                if i == 1:
+                    user_embedding = user_embedding.repeat(beam, 1)
+                    encoder_states = (encoder_states[0].repeat(beam, 1, 1),
+                                      encoder_states[1].repeat(beam, 1, 1))
+
+                scores, _ = self.decoder(xs.reshape(len(sequences[0])*bsz, -1), encoder_states)
+                scores = scores[:, -1:, :]
+                token_logits = F.linear(scores, self.token_embedding.weight)
+                user_logits = self.user_proj_2(torch.relu(self.user_proj_1(user_embedding))).unsqueeze(1)
+                sum_logits = token_logits + user_logits
+
+            logits = sum_logits.reshape(len(sequences[0]), bsz, 1, -1)
+            scores = scores.reshape(len(sequences[0]), bsz, 1, -1)
+            logits = torch.nn.functional.softmax(logits)  # turn into probabilities,in case of negative numbers
+            probs, preds = logits.topk(beam, dim=-1)
+            # (candeidate, bs, 1 , beam) during first loop, candidate=1, otherwise candidate=beam
+
+            for j in range(bsz):
+                all_candidates = []
+                for n in range(len(sequences[j])):
+                    for k in range(beam):
+                        prob = sequences[j][n][2]
+                        score = sequences[j][n][1]
+                        if score == []:
+                            score_tmp = scores[n][j][0].unsqueeze(0)
+                        else:
+                            score_tmp = torch.cat((score, scores[n][j][0].unsqueeze(0)), dim=0)
+                        seq_tmp = torch.cat((xs[n][j].reshape(-1), preds[n][j][0][k].reshape(-1)))
+                        candidate = [seq_tmp, score_tmp, prob * probs[n][j][0][k]]
+                        all_candidates.append(candidate)
+                ordered = sorted(all_candidates, key=lambda tup: tup[2], reverse=True)
+                sequences[j] = ordered[:beam]
+
+            # check if everyone has generated an end token
+            all_finished = ((xs == self.end_token_idx).sum(dim=1) > 0).sum().item() == bsz
+            if all_finished:
+                break
+        logits = torch.stack([seq[0][1] for seq in sequences])
+        xs = torch.stack([seq[0][0] for seq in sequences])
+        return logits, xs
+
     def converse(self, batch, mode):
         context_tokens, context_entities, response = batch['context_tokens'], batch['context_entities'], batch[
             'response']
@@ -240,3 +299,12 @@ def converse(self, batch, mode):
         else:
             _, preds = self.decode_greedy(encoder_state, user_embedding)
             return preds
+
+    def forward(self, batch, mode, stage):
+        if len(self.gpu) >= 2:
+            self.edge_idx = self.edge_idx.cuda(torch.cuda.current_device())
+            self.edge_type = self.edge_type.cuda(torch.cuda.current_device())
+        if stage == "conv":
+            return self.converse(batch, mode)
+        if stage == "rec":
+            return self.recommend(batch, mode)