Fixing tiny bugs

Author: codertimo

bert_pytorch/__main__.py

@@ -10,7 +10,7 @@
 def train():
     parser = argparse.ArgumentParser()
 
-    parser.add_argument("-d", "--train_dataset", required=True, type=str)
+    parser.add_argument("-c", "--train_dataset", required=True, type=str)
     parser.add_argument("-t", "--test_dataset", type=str, default=None)
     parser.add_argument("-v", "--vocab_path", required=True, type=str)
     parser.add_argument("-o", "--output_path", required=True, type=str)
@@ -23,7 +23,7 @@ def train():
     parser.add_argument("-b", "--batch_size", type=int, default=64)
     parser.add_argument("-e", "--epochs", type=int, default=10)
     parser.add_argument("-w", "--num_workers", type=int, default=5)
-    parser.add_argument("-c", "--with_cuda", type=bool, default=True)
+    parser.add_argument("--with_cuda", type=bool, default=True)
     parser.add_argument("--log_freq", type=int, default=10)
     parser.add_argument("--corpus_lines", type=int, default=None)
 

bert_pytorch/dataset/__init__.py

@@ -1,3 +1,2 @@
 from .dataset import BERTDataset
-from .creator import BERTDatasetCreator
 from .vocab import WordVocab

bert_pytorch/dataset/dataset.py

@@ -38,7 +38,7 @@ def __getitem__(self, item):
         output = {"bert_input": bert_input,
                   "bert_label": bert_label,
                   "segment_label": segment_label,
-                  "is_next": self.datas[item]["is_next"]}
+                  "is_next": is_next_label}
 
         return {key: torch.tensor(value) for key, value in output.items()}
 

bert_pytorch/dataset/vocab.py

@@ -167,7 +167,7 @@ def load_vocab(vocab_path: str) -> 'WordVocab':
             return pickle.load(f)
 
 
-if __name__ == "__main__":
+def build():
     import argparse
 
     parser = argparse.ArgumentParser()