Fix bug.

Author: xinyandai

README.md

@@ -15,9 +15,9 @@ string similarity search.
 python main.py --dataset word --nt 1000 --nq 1000 --epochs 20 --save-split --recall
 ```
 
-- set epoch as 0 to test bert embedding
+- test bert embedding
 ```
-python main.py --dataset word --nt 1000 --nq 1000 --epochs 0 --save-split --recall
+python main.py --dataset word --nt 1000 --nq 1000 --bert --save-split --recall
 ```
 
 ##### optional arguments:

embed_cnn.py

@@ -68,7 +68,7 @@ def cnn_embedding(args, h, data_file):
 
     # check if we use bert here
     char_alphabet = None
-    if args.epochs == 0:
+    if args.bert:
         char_alphabet = h.alphabet
 
     xt = _batch_embed(args, model.embedding_net, h.xt, device, char_alphabet=char_alphabet)

main.py

@@ -227,6 +227,9 @@ def get_args():
     parser.add_argument(
         "--recall", action="store_true", default=False, help="print recall"
     )
+    parser.add_argument(
+        "--bert", action="store_true", default=False, help="using bert or not"
+    )
     parser.add_argument("--embed", type=str, default="cnn", help="embedding method")
     parser.add_argument("--maxl", type=int, default=0, help="max length of strings")
     parser.add_argument(