diff --git a/runtime/components/stop_token_detector.cc b/runtime/components/stop_token_detector.cc index cb9f54e08..370fb327f 100644 --- a/runtime/components/stop_token_detector.cc +++ b/runtime/components/stop_token_detector.cc @@ -69,11 +69,8 @@ absl::Status StopTokenDetector::AddStopTokenSequence( } stop_sequences_storage_.push_back(stop_sequence); - - // Add a progress tracker for the new stop sequence for each batch item. - for (auto& progress_vector_for_item : batch_item_match_progress_) { - progress_vector_for_item.push_back(0); - } + max_stop_sequence_length_ = + std::max(max_stop_sequence_length_, stop_sequence.size()); return absl::OkStatus(); } @@ -81,10 +78,7 @@ void StopTokenDetector::ResetBatch(size_t batch_size) { int new_batch_size = batch_size == 0 ? stop_token_found_.size() : batch_size; stop_token_found_.assign(new_batch_size, false); max_batch_item_match_progress_.assign(new_batch_size, 0); - // Initialize progress for each batch item for all currently defined stop - // sequences. - batch_item_match_progress_.assign( - new_batch_size, std::vector(stop_sequences_storage_.size(), 0)); + batch_item_token_history_.assign(new_batch_size, std::vector()); matched_stop_sequence_length_.assign(new_batch_size, 0); } @@ -108,33 +102,51 @@ absl::Status StopTokenDetector::ProcessTokens( matched_stop_sequence_length_[i]++; continue; } - max_batch_item_match_progress_[i] = 0; + int current_token_id = latest_tokens[i]; - for (size_t k = 0; k < stop_sequences_storage_.size(); ++k) { - const auto& stop_seq_k = - stop_sequences_storage_[k]; // Guaranteed non-empty - int& current_match_len_for_k = batch_item_match_progress_[i][k]; - - if (current_match_len_for_k < stop_seq_k.size() && - stop_seq_k[current_match_len_for_k] == current_token_id) { - current_match_len_for_k++; - } else { - // Mismatch or sequence completed; reset progress for this stop_seq_k. - // Check if current token starts stop_seq_k anew. - if (stop_seq_k[0] == current_token_id) { - current_match_len_for_k = 1; - } else { - current_match_len_for_k = 0; + auto& history = batch_item_token_history_[i]; + history.push_back(current_token_id); + if (history.size() > max_stop_sequence_length_) { + history.erase(history.begin()); + } + + max_batch_item_match_progress_[i] = 0; + for (const auto& stop_seq : stop_sequences_storage_) { + // Check for full match. + if (history.size() >= stop_seq.size()) { + bool full_match = true; + for (size_t idx = 0; idx < stop_seq.size(); ++idx) { + if (history[history.size() - stop_seq.size() + idx] != + stop_seq[idx]) { + full_match = false; + break; + } + } + if (full_match) { + stop_token_found_[i] = true; + matched_stop_sequence_length_[i] = stop_seq.size(); + max_batch_item_match_progress_[i] = stop_seq.size(); + break; // Stop token found, no need to check other sequences. } } - if (current_match_len_for_k > 0 && - current_match_len_for_k == stop_seq_k.size()) { - stop_token_found_[i] = true; - matched_stop_sequence_length_[i] = stop_seq_k.size(); + // Check for partial match (only if stop token is not found yet). + if (!stop_token_found_[i]) { + for (int len = std::min(history.size(), stop_seq.size() - 1); + len > max_batch_item_match_progress_[i]; --len) { + bool partial_match = true; + for (int idx = 0; idx < len; ++idx) { + if (history[history.size() - len + idx] != stop_seq[idx]) { + partial_match = false; + break; + } + } + if (partial_match) { + max_batch_item_match_progress_[i] = len; + break; + } + } } - max_batch_item_match_progress_[i] = - std::max(max_batch_item_match_progress_[i], current_match_len_for_k); } } return absl::OkStatus(); diff --git a/runtime/components/stop_token_detector.h b/runtime/components/stop_token_detector.h index b01241bb1..7d1b38ab2 100644 --- a/runtime/components/stop_token_detector.h +++ b/runtime/components/stop_token_detector.h @@ -99,12 +99,16 @@ class StopTokenDetector { // Stores all added stop sequences. std::vector> stop_sequences_storage_; - // batch_item_match_progress_[i][k]: current match length for batch item 'i' - // against stop_sequences_storage_[k]. - std::vector> batch_item_match_progress_; + // Maximum length of any registered stop sequence. + size_t max_stop_sequence_length_ = 0; - // max_batch_item_match_progress_[i]: maximum match length for batch item 'i' - // against all stop_sequences_storage_[k]. + // The token history buffer for each batch item. + // batch_item_token_history_[i] stores the last max_stop_sequence_length_ + // tokens for batch item 'i'. + std::vector> batch_item_token_history_; + + // max_batch_item_match_progress_[i]: maximum partial match length for batch + // item 'i' against all stop_sequences_storage_[k]. std::vector max_batch_item_match_progress_; // stop_token_found_[i]: true if batch item 'i' has matched a stop sequence. diff --git a/runtime/components/stop_token_detector_test.cc b/runtime/components/stop_token_detector_test.cc index fc24241c2..b88904110 100644 --- a/runtime/components/stop_token_detector_test.cc +++ b/runtime/components/stop_token_detector_test.cc @@ -134,5 +134,27 @@ TEST(StopTokenDetectorTest, ResetBatch) { EXPECT_EQ(0, detector.GetStepsBeforeStopTokens()[0]); } +TEST(StopTokenDetectorTest, ProcessTokensSelfOverlap) { + StopTokenDetector detector(1); + EXPECT_OK(detector.AddStopTokenSequence({1, 1, 2})); + + std::vector tokens = {1, 1, 1, 2}; + + size_t i; + for (i = 0; i < tokens.size(); ++i) { + std::vector current_batch_tokens = {tokens[i]}; + EXPECT_OK(detector.ProcessTokens(absl::MakeSpan(current_batch_tokens))); + auto all_done = detector.AllDone(); + ASSERT_OK(all_done.status()); + if (*all_done) { + break; + } + } + EXPECT_EQ(i, 3); + auto final_all_done = detector.AllDone(); + ASSERT_OK(final_all_done.status()); + EXPECT_TRUE(*final_all_done); +} + } // namespace } // namespace litert::lm