From d26723f747cafa8d042f97033821773176b32c06 Mon Sep 17 00:00:00 2001
From: Anastasiia Pnevskaia <anastasiia.pnevskaia@intel.com>
Date: Fri, 12 Jul 2024 15:38:28 +0200
Subject: [PATCH] Clear beam search info.

---
 src/cpp/src/continuous_batching_pipeline.cpp | 1 +
 src/cpp/src/sampler.hpp                      | 6 ++++++
 2 files changed, 7 insertions(+)

diff --git a/src/cpp/src/continuous_batching_pipeline.cpp b/src/cpp/src/continuous_batching_pipeline.cpp
index dbacf3c243..beeabf4d01 100644
--- 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
index 5dc44b491f..cbc48a995b 100644
--- 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();
+}
 }