Clear beam search info.

Wovchena · Jul 12, 2024 · d26723f · d26723f
1 parent 165606a
commit d26723f
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/cpp/src/continuous_batching_pipeline.cpp b/src/cpp/src/continuous_batching_pipeline.cpp
@@ -269,6 +269,7 @@ class ContinuousBatchingPipeline::Impl {
             result.m_status = generation->get_status();
             results.push_back(result);
         }
+        m_sampler->clear_beam_search_info();
 
         OPENVINO_ASSERT(results.size() == prompts.size());
         return results;

diff --git a/src/cpp/src/sampler.hpp b/src/cpp/src/sampler.hpp
@@ -247,6 +247,8 @@ class Sampler {
     SamplerOutput sample(std::vector<SequenceGroup::Ptr> & sequence_groups, ov::Tensor logits);
 
     void set_seed(size_t seed) { rng_engine.seed(seed); }
+
+    void clear_beam_search_info();
 };
 
 SamplerOutput Sampler::sample(std::vector<SequenceGroup::Ptr> & sequence_groups, ov::Tensor logits) {
@@ -578,4 +580,8 @@ void GroupBeamSearcher::select_next_tokens(const ov::Tensor& logits, SamplerOutp
         }
     }
 }
+
+void Sampler::clear_beam_search_info() { 
+    m_beam_search_info.clear();
+}
 }