Merge remote-tracking branch 'popovaan/clear_beam_info' into use-CB-a…

…s-backend
Wovchena · Jul 12, 2024 · c28a023 · c28a023
2 parents 67f4717 + d26723f
commit c28a023
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
@@ -272,6 +272,7 @@ class ContinuousBatchingPipeline::Impl {
             result.m_status = generation->get_status();
             results.push_back(std::move(result));
         }
+        m_sampler->clear_beam_search_info();
 
         OPENVINO_ASSERT(results.size() == input_ids.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();
+}
 }