Add optional beam_width parameter to examples/beam_search/beam_search.cpp.

Author: mattpulver

examples/beam_search/beam_search.cpp

@@ -32,14 +32,15 @@
 int main(int argc, char ** argv)
 {
     gpt_params params;
+    //params.n_gpu_layers = 200;
 
     //---------------------------------
     // Print help :
     //---------------------------------
 
     if ( argc < 2 || argv[1][0] == '-' )
     {
-        printf( "Usage: %s MODEL_PATH [PROMPT]\n" , argv[0] );
+        printf( "Usage: %s MODEL_PATH [BEAM_WIDTH=2] [PROMPT]\n" , argv[0] );
         return 1 ;
     }
 
@@ -49,11 +50,11 @@ int main(int argc, char ** argv)
 
     params.model = argv[1];
 
-    params.n_beams = 2;  // Hard-code 2 until we can calculate how much memory is required
+    params.n_beams = 2 < argc ? std::stoi(argv[2]) : 2;
 
-    if ( argc > 2 )
+    if ( argc > 3 )
     {
-        params.prompt = argv[2];
+        params.prompt = argv[3];
     }
 
     if ( params.prompt.empty() )