Class T2MTextEncoder

diffplanner/models/rnns/t2m_bigru.py:106–161

@SUBMODULES.register_module()
class T2MTextEncoder(nn.Module):
    
    def __init__(self,
                 word_size,
                 pos_size,
                 hidden_size,
                 output_size,
                 max_text_len):
        super().__init__()
        self.text_encoder = TextEncoderBiGRUCo(
            word_size=word_size,
            pos_size=pos_size,
            hidden_size=hidden_size,
            output_size=output_size,
        )
        self.w_vectorizer = WordVectorizer('./data/glove', 'our_vab')
        self.max_text_len = max_text_len
        
    def load_pretrained(self, ckpt_path):
        checkpoint = torch.load(ckpt_path, map_location='cpu')
        self.text_encoder.load_state_dict(checkpoint['text_encoder'])
        
    def forward(self, text, token, device):
        B = len(text)
        pos_one_hot = []
        word_emb = []
        sent_len = []
        for i in range(B):
            tokens = token[i].split(" ")
            if len(tokens) < self.max_text_len:
                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
                batch_sent_len = len(tokens)
                tokens = tokens + ['unk/OTHER'] * (self.max_text_len + 2 - batch_sent_len)
            else:
                tokens = tokens[: self.max_text_len]
                tokens = ['sos/OTHER'] + tokens + ['eos/OTHER']
                batch_sent_len = len(tokens)
            sent_len.append(batch_sent_len)
            batch_word_emb = []
            batch_pos_one_hot = []
            for cur_token in tokens:
                cur_word_emb, cur_pos_one_hot = self.w_vectorizer[cur_token]
                cur_word_emb = torch.from_numpy(cur_word_emb).float()
                cur_pos_one_hot = torch.from_numpy(cur_pos_one_hot).float()
                batch_word_emb.append(cur_word_emb)
                batch_pos_one_hot.append(cur_pos_one_hot)
            
            batch_word_emb = torch.stack(batch_word_emb, dim=0)
            batch_pos_one_hot = torch.stack(batch_pos_one_hot, dim=0)
            word_emb.append(batch_word_emb)
            pos_one_hot.append(batch_pos_one_hot)
        word_emb = torch.stack(word_emb, dim=0).to(device)
        pos_one_hot = torch.stack(pos_one_hot, dim=0).to(device)
        sent_len = torch.tensor(sent_len, dtype=torch.long).to(device)
        text_embedding = self.text_encoder(word_emb, pos_one_hot, sent_len)
        return text_embedding