| 89 | return sents |
| 90 | |
| 91 | class BatchSentLoader(object): |
| 92 | def __init__(self, sents, batch_size, pad_id=0, cuda=False, volatile=False): |
| 93 | self.sents = sents |
| 94 | self.batch_size = batch_size |
| 95 | self.sort_sents = sorted(sents, key=lambda x: x.size(0)) |
| 96 | self.cuda = cuda |
| 97 | self.volatile = volatile |
| 98 | self.pad_id = pad_id |
| 99 | |
| 100 | def __next__(self): |
| 101 | if self.idx >= len(self.sort_sents): |
| 102 | raise StopIteration |
| 103 | |
| 104 | batch_size = min(self.batch_size, len(self.sort_sents)-self.idx) |
| 105 | batch = self.sort_sents[self.idx:self.idx+batch_size] |
| 106 | max_len = max([s.size(0) for s in batch]) |
| 107 | tensor = torch.LongTensor(max_len, batch_size).fill_(self.pad_id) |
| 108 | for i in range(len(batch)): |
| 109 | s = batch[i] |
| 110 | tensor[:s.size(0),i].copy_(s) |
| 111 | if self.cuda: |
| 112 | tensor = tensor.cuda() |
| 113 | |
| 114 | self.idx += batch_size |
| 115 | |
| 116 | return tensor |
| 117 | |
| 118 | next = __next__ |
| 119 | |
| 120 | def __iter__(self): |
| 121 | self.idx = 0 |
| 122 | return self |
| 123 | |
| 124 | if __name__ == '__main__': |
| 125 | corpus = SentCorpus('../penn') |