From eb0b38ea82344a7edc5f768dac42ed4905ee0f95 Mon Sep 17 00:00:00 2001
From: Rudolf A Braun <ruseni@amazon.com>
Date: Mon, 15 Sep 2025 12:07:50 -0400
Subject: [PATCH] rewind

---
 bindings/python/flashlight/lib/text/_decoder.cpp | 1 +
 flashlight/lib/text/decoder/LexiconDecoder.cpp   | 9 +++++++++
 flashlight/lib/text/decoder/LexiconDecoder.h     | 2 ++
 3 files changed, 12 insertions(+)

diff --git a/bindings/python/flashlight/lib/text/_decoder.cpp b/bindings/python/flashlight/lib/text/_decoder.cpp
index 177bedc5..0eaec0df 100644
--- a/bindings/python/flashlight/lib/text/_decoder.cpp
+++ b/bindings/python/flashlight/lib/text/_decoder.cpp
@@ -363,6 +363,7 @@ PYBIND11_MODULE(flashlight_lib_text_decoder, m) {
       .def("decode_end", &LexiconDecoder::decodeEnd)
       .def("decode", &LexiconDecoder_decode, "emissions"_a, "T"_a, "N"_a)
       .def("prune", &LexiconDecoder::prune, "look_back"_a = 0)
+      .def("rewind", &LexiconDecoder::rewind, "num_frames"_a)
       .def(
           "get_best_hypothesis",
           &LexiconDecoder::getBestHypothesis,
diff --git a/flashlight/lib/text/decoder/LexiconDecoder.cpp b/flashlight/lib/text/decoder/LexiconDecoder.cpp
index 0590c088..f08c4fd0 100644
--- a/flashlight/lib/text/decoder/LexiconDecoder.cpp
+++ b/flashlight/lib/text/decoder/LexiconDecoder.cpp
@@ -323,6 +323,15 @@ void LexiconDecoder::prune(int lookBack) {
 
   nPrunedFrames_ = nDecodedFrames_ - lookBack;
 }
+
+void LexiconDecoder::rewind(int numFrames) {
+  int startFrame = std::max(1, nDecodedFrames_ - numFrames);
+  for (int i = startFrame; i <= nDecodedFrames_; i++) {
+    hyp_.erase(i);
+  }
+
+  nDecodedFrames_ = std::max(0, nDecodedFrames_ - numFrames);
+}
 } // namespace text
 } // namespace lib
 } // namespace fl
diff --git a/flashlight/lib/text/decoder/LexiconDecoder.h b/flashlight/lib/text/decoder/LexiconDecoder.h
index 019dc94a..0d39227d 100644
--- a/flashlight/lib/text/decoder/LexiconDecoder.h
+++ b/flashlight/lib/text/decoder/LexiconDecoder.h
@@ -142,6 +142,8 @@ class FL_TEXT_API LexiconDecoder : public Decoder {
 
   void prune(int lookBack = 0) override;
 
+  void rewind(int numFrames);
+
   int nDecodedFramesInBuffer() const override;
 
   DecodeResult getBestHypothesis(int lookBack = 0) const override;