INNER CODE UNIT · Python
token_ids
yxuansu/SimCTG · dialogue_generation/dataclass.py:79
token_ids = token_ids[-self.max_len:]
assert len(token_ids) <= self.max_len
res_token_list.append(tokens)
res_token_id_list.append(token_ids)
return
def pad_batch(self, batch_id_list):
batch_id_list = [torch.LongTensor(item) for item in batch_id_list]
batch_tensor = rnn.pad_sequence(batch_id_list, batch_first=True, padding_value=self.pad_token_id)
batch_mask = torch.ones_like(batch_tensor)
batch_mask = batch_mask.masked_fill(batch_tensor.eq(self.pad_token_id), 0.0).type(torch.FloatTensor)
return batch_tensor, batch_mask
def process_output(self, batch_tgt_id_list):
batch_tgt_id_list = [torch.LongTensor(item) for item in batch_tgt_id_list]
batch_tgt_tensor, _ = self.pad_batch(batch_tgt_id_list) # padded target sequence
batch_tgt_input_tensor = batch_tgt_tensor[:, :-1].clone()
batch_tgt_output_tensor = batch_tgt_tensor[:, 1:].clone()