if

flexflow · Dec 16, 2024 · 959ec78 · 959ec78
1 parent 57739b6
commit 959ec78
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/inference/models/mixtral.cc b/inference/models/mixtral.cc
@@ -383,15 +383,23 @@ void MIXTRAL::create_mixtral_model(FFModel &ff,
                           "lm_head");
 
   Tensor output;
+    if (mode == BEAM_SEARCH_MODE) {
+    Tensor softmax = ff.softmax(dense, -1);
+    // output = ff.beam_top_k(softmax, mixtral_config.max_beam_width, false);
+    // output = ff.argmax(softmax, /*beam_Search*/ true);
+    output = ff.arg_top_k(softmax, mixtral_config.max_beam_width, false, true);
+    // output = ff.top_k(softmax, )
+  } else {
     // Tensor softmax = ff.softmax(dense, -1);
     if (generation_config.do_sample) {
       dense = ff.scalar_truediv(dense, generation_config.temperature, false);
       Tensor softmax = ff.softmax(dense, -1);
       output = ff.sampling(softmax, generation_config.topp);
-    } else {
+   } else {
       Tensor softmax = ff.softmax(dense, -1);
       output = ff.argmax(softmax, /*beam_Search*/ false);
     }
+  }
 
   FileDataLoader *fileloader = new FileDataLoader(
       "",