Merge pull request #366 from donglihe-hub/beam_width

add beam_width

Author: Eleven1Liu

main.py

@@ -223,6 +223,12 @@ def add_all_arguments(parser):
     parser.add_argument(
         "--tree_max_depth", type=int, default=10, help="Maximum depth of the tree (default: %(default)s)"
     )
+    parser.add_argument(
+        "--beam_width",
+        type=int,
+        default=10,
+        help="The width of the beam search (default: %(default)s)",
+    )
     parser.add_argument(
         "-h",
         "--help",

tests/basic.sh

@@ -119,7 +119,7 @@ find example_config -name "*.yml" ${filters[@]} -type f -print0 |
         echo "Running $config"
         stderr=$(python $script --config "$config" --epochs 1 \
             --result_dir "$result_dir" --embed_cache_dir data \
-            --val_size 0.2 --save_k_predictions 2 \
+            --val_size 0.2 --beam_width 1 \
             --cpu 2>&1 > /dev/null)
         if [[ $? -ne 0 ]]; then
             echo "$stderr" >&2