| 56 | return ids |
| 57 | |
| 58 | class SentCorpus(object): |
| 59 | def __init__(self, path): |
| 60 | self.dictionary = Dictionary() |
| 61 | self.train = self.tokenize(os.path.join(path, 'train.txt')) |
| 62 | self.valid = self.tokenize(os.path.join(path, 'valid.txt')) |
| 63 | self.test = self.tokenize(os.path.join(path, 'test.txt')) |
| 64 | |
| 65 | def tokenize(self, path): |
| 66 | """Tokenizes a text file.""" |
| 67 | assert os.path.exists(path) |
| 68 | # Add words to the dictionary |
| 69 | with open(path, 'r', encoding='utf-8') as f: |
| 70 | tokens = 0 |
| 71 | for line in f: |
| 72 | words = line.split() + ['<eos>'] |
| 73 | tokens += len(words) |
| 74 | for word in words: |
| 75 | self.dictionary.add_word(word) |
| 76 | |
| 77 | # Tokenize file content |
| 78 | sents = [] |
| 79 | with open(path, 'r', encoding='utf-8') as f: |
| 80 | for line in f: |
| 81 | if not line: |
| 82 | continue |
| 83 | words = line.split() + ['<eos>'] |
| 84 | sent = torch.LongTensor(len(words)) |
| 85 | for i, word in enumerate(words): |
| 86 | sent[i] = self.dictionary.word2idx[word] |
| 87 | sents.append(sent) |
| 88 | |
| 89 | return sents |
| 90 | |
| 91 | class BatchSentLoader(object): |
| 92 | def __init__(self, sents, batch_size, pad_id=0, cuda=False, volatile=False): |