beam_search_callback_state struct can be heavy and referenced in callback.

Author: mattpulver

examples/beam_search/beam_search.cpp

@@ -43,7 +43,7 @@ std::ostream& operator<<(std::ostream& os, ostream_beam_view const& obv) {
 // Put here anything you want back in beam_search_callback().
 struct beam_search_callback_state {
     llama_context* ctx;
-    std::vector<llama_token>* response;
+    std::vector<llama_token> response;
 };
 
 bool is_at_eos(beam_search_callback_state, llama_token const* tokens, size_t const n_tokens) {
@@ -56,28 +56,28 @@ bool is_at_eos(beam_search_callback_state, llama_token const* tokens, size_t con
 //  * When all beams converge to a common prefix, they are made available in beams_state.beams[0].
 //    This is also called when the stop condition is met.
 //    Collect tokens into std::vector<llama_token> response which is pointed to by callback_state.
-void beam_search_callback(void* callback_state, llama_beams_state beams_state) {
-    auto const state = *static_cast<beam_search_callback_state*>(callback_state);
+void beam_search_callback(void* callback_state_ptr, llama_beams_state beams_state) {
+    auto& callback_state = *static_cast<beam_search_callback_state*>(callback_state_ptr);
     // Mark beams as EOS as needed.
     for (size_t i=0 ; i<beams_state.n_beams ; ++i) {
         llama_beam_view& beam_view = beams_state.beam_views[i];
-        if (!beam_view.eos && is_at_eos(state, beam_view.tokens, beam_view.n_tokens)) {
+        if (!beam_view.eos && is_at_eos(callback_state, beam_view.tokens, beam_view.n_tokens)) {
             beam_view.eos = true;
         }
     }
     printf(",");  // Show progress
     if (size_t const n = beams_state.common_prefix_length) {
-        state.response->resize(state.response->size() + n);
+        callback_state.response.resize(callback_state.response.size() + n);
         assert(0u < beams_state.n_beams);
         llama_token const* tokens = beams_state.beam_views[0].tokens;
-        std::copy(tokens, tokens + n, state.response->end() - n);
+        std::copy(tokens, tokens + n, callback_state.response.end() - n);
         printf("%lu", n);
     }
     fflush(stdout);
 #if 1 // DEBUG: print current beams for this iteration
-    std::cout << "\n\nCurrent beams:\n";
+    std::cout << "\n\nCurrent beams (last_call=" << beams_state.last_call << "):\n";
     for (size_t i=0 ; i < beams_state.n_beams ; ++i) {
-        std::cout << "beams["<<i<<"]: " << ostream_beam_view{state.ctx,beams_state.beam_views[i]} << std::endl;
+        std::cout << "beams["<<i<<"]: " << ostream_beam_view{callback_state.ctx,beams_state.beam_views[i]} << std::endl;
     }
 #endif
 }
@@ -168,14 +168,13 @@ int main(int argc, char ** argv)
     }
     n_past += tokens_list.size();
 
-    std::vector<llama_token> response;
-    beam_search_callback_state callback_state{ctx, &response};
+    beam_search_callback_state callback_state{ctx, {}};
     size_t const beam_width = static_cast<size_t>(params.n_beams);
     int const n_predict = 256;
     llama_beam_search(ctx, beam_search_callback, &callback_state, beam_width, n_past, n_predict, params.n_threads);
 
     printf("\n\n");
-    for (llama_token const token_id : response) {
+    for (llama_token const token_id : callback_state.response) {
         printf("%s", llama_token_to_str(ctx,token_id));
     }
     printf("\n");