init

lib/datasets/LanguageDataset.py

@@ -0,0 +1,122 @@
+import os
+import torch
+
+from collections import Counter
+
+
+class Dictionary(object):
+  def __init__(self):
+    self.word2idx = {}
+    self.idx2word = []
+    self.counter = Counter()
+    self.total = 0
+
+  def add_word(self, word):
+    if word not in self.word2idx:
+      self.idx2word.append(word)
+      self.word2idx[word] = len(self.idx2word) - 1
+    token_id = self.word2idx[word]
+    self.counter[token_id] += 1
+    self.total += 1
+    return self.word2idx[word]
+
+  def __len__(self):
+    return len(self.idx2word)
+
+
+class Corpus(object):
+  def __init__(self, path):
+    self.dictionary = Dictionary()
+    self.train = self.tokenize(os.path.join(path, 'train.txt'))
+    self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
+    self.test = self.tokenize(os.path.join(path, 'test.txt'))
+
+  def tokenize(self, path):
+    """Tokenizes a text file."""
+    assert os.path.exists(path)
+    # Add words to the dictionary
+    with open(path, 'r', encoding='utf-8') as f:
+      tokens = 0
+      for line in f:
+        words = line.split() + ['<eos>']
+        tokens += len(words)
+        for word in words:
+          self.dictionary.add_word(word)
+
+    # Tokenize file content
+    with open(path, 'r', encoding='utf-8') as f:
+      ids = torch.LongTensor(tokens)
+      token = 0
+      for line in f:
+        words = line.split() + ['<eos>']
+        for word in words:
+          ids[token] = self.dictionary.word2idx[word]
+          token += 1
+
+    return ids
+
+class SentCorpus(object):
+  def __init__(self, path):
+    self.dictionary = Dictionary()
+    self.train = self.tokenize(os.path.join(path, 'train.txt'))
+    self.valid = self.tokenize(os.path.join(path, 'valid.txt'))
+    self.test = self.tokenize(os.path.join(path, 'test.txt'))
+
+  def tokenize(self, path):
+    """Tokenizes a text file."""
+    assert os.path.exists(path)
+    # Add words to the dictionary
+    with open(path, 'r', encoding='utf-8') as f:
+      tokens = 0
+      for line in f:
+        words = line.split() + ['<eos>']
+        tokens += len(words)
+        for word in words:
+          self.dictionary.add_word(word)
+
+    # Tokenize file content
+    sents = []
+    with open(path, 'r', encoding='utf-8') as f:
+      for line in f:
+        if not line:
+          continue
+        words = line.split() + ['<eos>']
+        sent = torch.LongTensor(len(words))
+        for i, word in enumerate(words):
+          sent[i] = self.dictionary.word2idx[word]
+        sents.append(sent)
+
+    return sents
+
+class BatchSentLoader(object):
+  def __init__(self, sents, batch_size, pad_id=0, cuda=False, volatile=False):
+    self.sents = sents
+    self.batch_size = batch_size
+    self.sort_sents = sorted(sents, key=lambda x: x.size(0))
+    self.cuda = cuda
+    self.volatile = volatile
+    self.pad_id = pad_id
+
+  def __next__(self):
+    if self.idx >= len(self.sort_sents):
+      raise StopIteration
+
+    batch_size = min(self.batch_size, len(self.sort_sents)-self.idx)
+    batch = self.sort_sents[self.idx:self.idx+batch_size]
+    max_len = max([s.size(0) for s in batch])
+    tensor = torch.LongTensor(max_len, batch_size).fill_(self.pad_id)
+    for i in range(len(batch)):
+      s = batch[i]
+      tensor[:s.size(0),i].copy_(s)
+    if self.cuda:
+      tensor = tensor.cuda()
+
+    self.idx += batch_size
+
+    return tensor
+  
+  next = __next__
+
+  def __iter__(self):
+    self.idx = 0
+    return self