Merge pull request #23 from bazingagin/exp

fix custom dataset

README.md

@@ -47,7 +47,9 @@ to calculate accuracy. Otherwise, the accuracy will be calculated automatically
 
 ### Use Custom Dataset
 
-You can use your own custom dataset by passing whatever string to `--dataset`, and also remember to pass the data directory that contains `train.txt` and `test.txt` to `--data_dir`.
+You can use your own custom dataset by passing `custom` to `--dataset`; pass the data directory that contains `train.txt` and `test.txt` to `--data_dir`; pass the class number to the `--class_num`.
+
 Both `train.txt` and `test.txt` are expected to have the format `{label}\t{text}` per line.
+
 You can change the delimiter according to you dataset by changing `delimiter` in `load_custom_dataset()` in `data.py`.
 

main_text.py

@@ -118,6 +118,7 @@ def non_neurl_knn_exp_given_dis(dis_matrix, k, test_label, train_label):
     parser.add_argument("--distance_fn", default=None)
     parser.add_argument("--score", action="store_true", default=False)
     parser.add_argument("--k", default=2, type=int)
+    parser.add_argument("--class_num", default=5, type=int)
     args = parser.parse_args()
     # create output dir
     if not os.path.exists(args.output_dir):
@@ -145,6 +146,7 @@ def non_neurl_knn_exp_given_dis(dis_matrix, k, test_label, train_label):
         "swahili": 6,
         "filipino": 5,
         "kirnews": 14,
+        "custom": args.class_num
     }
     # load dataset
     data_dir = os.path.join(args.data_dir, args.dataset)
@@ -158,6 +160,7 @@ def non_neurl_knn_exp_given_dis(dis_matrix, k, test_label, train_label):
         "swahili",
         "filipino",
         "kirnews",
+        "custom",
     ]:
         dataset_pair = eval(args.dataset)(root=args.data_dir)
     else: