Commit a7d0bd0e authored by alexeib's avatar alexeib Committed by Myle Ott
Browse files

fix token block rotation

parent 0e9e7f7b

fairseq/data/token_block_dataset.py

+8 −4
Changes for fairseq/data/token_block_dataset.py: 8 added lines, 4 removed lines.
Original line number Diff line number Diff line
@@ -76,13 +76,17 @@ class TokenBlockDataset(torch.utils.data.Dataset): 


    def __getitem__(self, index):

        s, e = self.slice_indices[index]



        item = torch.LongTensor(self.tokens[s:e])

        if self.include_targets:

            if e == self.total_size:

                return item[:-1], item[1:]

            else:

                return item, torch.LongTensor(self.tokens[s + 1:e + 1])

            # target is the sentence, for source, rotate item one token to the left (would start with eos)

            if s == 0:

                source = np.concatenate([self.tokens[-1:], self.tokens[0:e - 1]])

            else:

                source = self.tokens[s - 1:e - 1]



            return torch.LongTensor(source), item



        return item



    def __len__(self):