fix bugs for quesnet

Author: wintermelon008

EduNLP/ModelZoo/quesnet/quesnet.py

@@ -114,6 +114,10 @@ def make_batch(self, data, device, pretrain=False):
         ans_input = []
         ans_output = []
         false_options = [[] for i in range(3)]
+
+        if not isinstance(data, list):
+            data = [data]
+
         for q in data:
             meta = torch.zeros(len(self.stoi[self.meta])).to(device)
             meta[q.labels.get(self.meta) or []] = 1
@@ -156,7 +160,7 @@ def make_batch(self, data, device, pretrain=False):
 
             for i, fo in enumerate(q.false_options):
                 false_options[i].append([0] + fo)
-
+     
         lembs = SeqBatch(lembs, device=device)
         rembs = SeqBatch(rembs, device=device)
         embs = SeqBatch(embs, device=device)
@@ -192,6 +196,23 @@ def make_batch(self, data, device, pretrain=False):
         words = torch.cat(words, dim=0) if words else None
         ims = torch.cat(ims, dim=0) if ims else None
         metas = torch.cat(metas, dim=0) if metas else None
+        
+        
+        # print("debug1")
+        # print(lembs)
+        # print(rembs)
+        # print(words)
+        # print(ims)
+        # print(metas)
+        # print(wmask)
+        # print(imask)
+        # print(mmask)
+        # print(embs)
+        # print(ans_input)
+        # print(ans_output)
+        # print(false_opt_input)
+        
+        
         if pretrain:
             return (
                 lembs, rembs, words, ims, metas, wmask, imask, mmask,
@@ -302,67 +323,70 @@ def __init__(self, _stoi=None, pretrained_embs: np.ndarray = None, pretrained_im
         self.config = PretrainedConfig.from_dict(self.config)
 
     def forward(self, batch):
-        left, right, words, ims, metas, wmask, imask, mmask, inputs, ans_input, ans_output, false_opt_input = batch
+        left, right, words, ims, metas, wmask, imask, mmask, inputs, ans_input, ans_output, false_opt_input = batch[0]
 
         # high-level loss
         outputs = self.quesnet(inputs)
         embeded = outputs.embeded
         h = outputs.hidden
 
         x = ans_input.packed()
-        y, _ = self.ans_decode(PackedSequence(self.quesnet.we(x.data), x.batch_sizes),
+        
+        y, _ = self.ans_decode(PackedSequence(self.quesnet.we(x[0].data), x.batch_sizes),
                                h.repeat(self.config.layers, 1, 1))
         floss = F.cross_entropy(self.ans_output(y.data),
                                 ans_output.packed().data)
         floss = floss + F.binary_cross_entropy_with_logits(self.ans_judge(y.data),
                                                            torch.ones_like(self.ans_judge(y.data)))
         for false_opt in false_opt_input:
             x = false_opt.packed()
-            y, _ = self.ans_decode(PackedSequence(self.quesnet.we(x.data), x.batch_sizes),
+            if x == (None, None):
+                continue
+            y, _ = self.ans_decode(PackedSequence(self.quesnet.we(x[0].data), x.batch_sizes),
                                    h.repeat(self.config.layers, 1, 1))
             floss = floss + F.binary_cross_entropy_with_logits(self.ans_judge(y.data),
                                                                torch.zeros_like(self.ans_judge(y.data)))
         loss = floss * self.lambda_loss[1]
         # low-level loss
-        left_hid = self.quesnet(left).pack_embeded.data[:, :self.rnn_size]
-        right_hid = self.quesnet(right).pack_embeded.data[:, self.rnn_size:]
+        left_hid = self.quesnet(left).pack_embeded.data[:, :self.rnn_size].clone()
+        right_hid = self.quesnet(right).pack_embeded.data[:, self.rnn_size:].clone()
 
         wloss = iloss = mloss = None
 
         if words is not None:
-            lwfea = torch.masked_select(left_hid, wmask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            lout = self.lwoutput(lwfea)
-            rwfea = torch.masked_select(right_hid, wmask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            rout = self.rwoutput(rwfea)
-            out = self.woutput(torch.cat([lwfea, rwfea], dim=1))
+            lwfea = torch.masked_select(left_hid.clone(), wmask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            lout = self.lwoutput(lwfea.clone())
+            rwfea = torch.masked_select(right_hid.clone(), wmask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            rout = self.rwoutput(rwfea.clone())
+            out = self.woutput(torch.cat([lwfea.clone(), rwfea.clone()], dim=1).clone())
             wloss = (F.cross_entropy(out, words) + F.cross_entropy(lout, words) + F.
                      cross_entropy(rout, words)) * self.quesnet.lambda_input[0] / 3
             wloss *= self.lambda_loss[0]
             loss = loss + wloss
 
         if ims is not None:
-            lifea = torch.masked_select(left_hid, imask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            lout = self.lioutput(lifea)
-            rifea = torch.masked_select(right_hid, imask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            rout = self.rioutput(rifea)
-            out = self.ioutput(torch.cat([lifea, rifea], dim=1))
+            lifea = torch.masked_select(left_hid.clone(), imask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            lout = self.lioutput(lifea.clone())
+            rifea = torch.masked_select(right_hid.clone(), imask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            rout = self.rioutput(rifea.clone())
+            out = self.ioutput(torch.cat([lifea.clone(), rifea.clone()], dim=1).clone())
             iloss = (self.quesnet.ie.loss(ims, out) + self.quesnet.ie.loss(ims, lout) + self.quesnet.ie.
                      loss(ims, rout)) * self.quesnet.lambda_input[1] / 3
             iloss *= self.lambda_loss[0]
             loss = loss + iloss
 
         if metas is not None:
-            lmfea = torch.masked_select(left_hid, mmask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            lout = self.lmoutput(lmfea)
-            rmfea = torch.masked_select(right_hid, mmask.unsqueeze(1).bool()) \
-                .view(-1, self.rnn_size)
-            rout = self.rmoutput(rmfea)
-            out = self.moutput(torch.cat([lmfea, rmfea], dim=1))
+            lmfea = torch.masked_select(left_hid.clone(), mmask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            lout = self.lmoutput(lmfea.clone())
+            rmfea = torch.masked_select(right_hid.clone(), mmask.unsqueeze(1).bool()) \
+                .view(-1, self.rnn_size).clone()
+            rout = self.rmoutput(rmfea.clone())
+            out = self.moutput(torch.cat([lmfea.clone(), rmfea.clone()], dim=1).clone())
             mloss = (self.quesnet.me.loss(metas, out) + self.quesnet.me.loss(metas, lout) + self.quesnet.me.
                      loss(metas, rout)) * self.quesnet.lambda_input[2] / 3
             mloss *= self.lambda_loss[0]

EduNLP/ModelZoo/quesnet/util.py

@@ -11,14 +11,19 @@ def __init__(self, seqs, dtype=None, device=None):
         self.dtype = dtype
         self.device = device
         self.seqs = seqs
-        self.lens = [len(x) for x in seqs]
+
+        if not seqs:
+            self.lens = [0]
+        else:
+            self.lens = [len(x) for x in seqs]
 
         self.ind = argsort(self.lens)[::-1]
         self.inv = argsort(self.ind)
         self.lens.sort(reverse=True)
         self._prefix = [0]
         self._index = {}
         c = 0
+        
         for i in range(self.lens[0]):
             for j in range(len(self.lens)):
                 if self.lens[j] <= i:
@@ -28,10 +33,15 @@ def __init__(self, seqs, dtype=None, device=None):
 
     def packed(self):
         ind = torch.tensor(self.ind, dtype=torch.long, device=self.device)
+        if not ind.numel() or ind.max() >= self.padded()[0].size(1):
+            return None, None
         padded = self.padded()[0].index_select(1, ind)
         return pack_padded_sequence(padded, torch.tensor(self.lens))
 
     def padded(self, max_len=None, batch_first=False):
+        if not self.seqs:
+            return torch.empty((0, 0), dtype=self.dtype, device=self.device), torch.empty((0, 0), dtype=torch.bool, device=self.device)
+    
         seqs = [torch.tensor(s, dtype=self.dtype, device=self.device)
                 if not isinstance(s, torch.Tensor) else s
                 for s in self.seqs]

EduNLP/Pretrain/quesnet_vec.py

@@ -156,8 +156,9 @@ def _convert_to_ids(self, item: Union[str, dict, list], key=lambda x: x,
         token_item = self.tokenize(item, key)
         token_idx = []
         for _, w in enumerate(token_item):
-            if isinstance(w, FigureSegment):
+            if isinstance(w, FigureSegment) and 'ques_figure_ids' in item.keys():
                 # image                
+
                 try:
                     fig_id = f"{w.src[10:-1]}"
                     fig_index = item['ques_figure_ids'].index(fig_id)
@@ -171,11 +172,13 @@ def _convert_to_ids(self, item: Union[str, dict, list], key=lambda x: x,
                     else:
                         fig_src = item['ques_figure_paths'][fig_index]
 
+                    print(f"Open figure {fig_src}")
                     im = Image.open(fig_src)
                     im = im.resize((56, 56))
                     token_idx.append(to_grayscale(im))
+                    
                 except Exception:
-                    warnings.warn('Open image error! path = ' + fig_src)
+                    warnings.warn('Open image error!')
                     token_idx.append(self.stoi['word'][self.img_token])
             else:
                 # word
@@ -390,6 +393,7 @@ def __getitem__(self, index):
             token = self.tokenizer(line, key=self.content_key, meta=self.meta)
             content = token['seq_idx']
             meta = token['meta_idx']
+            
             if self.answer_key(line).isalpha() and len(self.answer_key(line)) == 1 and ord(self.answer_key(line)) < 128 and len(self.option_key(line)) > 0:
                 answer_idx = ord(self.answer_key(line).upper()) - ord('A')
                 options = self.option_key(line)
@@ -441,7 +445,7 @@ def __init__(self, data, *label, length=None, batch_size=1, shuffle=True):
         self.batch_size = batch_size
         self.queue = queue.Queue(maxsize=8)
         self.length = length if length is not None else len(data)
-
+        
         assert all(self.length == len(lab) for lab in label), \
             'data and label must have same lengths'
 
@@ -545,7 +549,7 @@ def optimizer(*models, **kwargs):
         return _cur_optim   
 
 
-def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save_embs = False, train_params = None):
+def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save_embs = False, load_embs = False, train_params = None):
     """ pretrain quesnet
 
     Parameters
@@ -558,6 +562,8 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
         quesnet tokenizer
     save_embs : bool, optional
         whether to save pretrained word/image/meta embeddings seperately
+    load_embs : bool, optional
+        whether to load pretrained word/image/meta embeddings seperately
     train_params : dict, optional
         the training parameters and model parameters, by default None
         - "n_epochs": int, default = 1
@@ -609,7 +615,7 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
         default_train_params.update(train_params)
     train_params = default_train_params
 
-    dataset = QuesnetDataset(path)
+    dataset = QuesnetDataset(path, img_dir=img_dir)
     tokenizer = dataset.tokenizer
     tokenizer.save_pretrained(output_dir)
     model = QuesNetForPreTraining(_stoi=tokenizer.stoi, feat_size=train_params['feat_size'],
@@ -642,7 +648,7 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
         meta_corpus.append(meta_vector)
 
     # train word2vec for text embedding
-    if pretrain_dir != None:
+    if pretrain_dir != None and load_embs:
         model.quesnet.load_emb(np.load(os.path.join(output_dir, 'w2v_embs.npy')))
     else:
         gensim_w2v = Word2Vec(sentences=[[item] for item in emb_dict.keys()], min_count=1,
@@ -661,7 +667,7 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
     logger.info('quesnet Word Embedding loaded')
 
     # train auto-encoder loss for image embedding
-    if pretrain_dir != None:
+    if pretrain_dir != None and load_embs:
         model.quesnet.load_img(torch.load(os.path.join(pretrain_dir, 'trained_ie.pt')))
     else:
         img_dataset = EmbeddingDataset(data=img_corpus, data_type='image')
@@ -675,7 +681,7 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
 
 
     # train auto-encoder loss for meta embedding
-    if pretrain_dir != None:
+    if pretrain_dir != None and load_embs:
         model.quesnet.load_meta(torch.load(os.path.join(pretrain_dir, 'trained_me.pt')))
     else:
         meta_dateset = EmbeddingDataset(data=meta_corpus, data_type='meta')
@@ -696,7 +702,7 @@ def pretrain_quesnet(path, output_dir, pretrain_dir = None, img_dir = None, save
     optim = optimizer(model, lr=train_params['lr'])
     n_batches = 0
     for epoch in range(0, train_params['n_epochs']):
-        train_iter = PrefetchIter(dataset, train_params['batch_size'])
+        train_iter = PrefetchIter(dataset, batch_size=train_params['batch_size'])
         bar = enumerate(tqdm(train_iter, initial=train_iter.pos),
                         train_iter.pos)
         for i, batch in critical(bar):