Fix issue of outputting stop token in Decode/DecodeStreaming

LiteRT-LM-PiperOrigin-RevId: 789000868

Author: ai-edge-bot

runtime/core/pipeline.cc

@@ -341,9 +341,11 @@ absl::StatusOr<Responses> Decode(LlmExecutor& executor, Tokenizer& tokenizer,
     }
 
     previous_token_ids.resize(0);
-    response_texts[0] +=
-        absl::StrReplaceAll(run_one_step.GetResultTokens()[0], {{"▁", " "}});
     num_decoded_steps++;
+    if (decode_result == kPartial) {
+      response_texts[0] +=
+          absl::StrReplaceAll(run_one_step.GetResultTokens()[0], {{"▁", " "}});
+    }
 
     if (ShouldStop(decode_result == kDone, benchmark_decode_token_count,
                    num_decoded_steps, executor.GetCurrentStep().value(),
@@ -398,10 +400,12 @@ absl::Status DecodeStreaming(LlmExecutor& executor, Tokenizer& tokenizer,
     }
 
     previous_token_ids.resize(0);
-    response_texts[0] +=
-        absl::StrReplaceAll(run_one_step.GetResultTokens()[0], {{"▁", " "}});
     num_decoded_steps++;
-    observer->OnNext(responses);
+    if (decode_result == kPartial) {
+      response_texts[0] +=
+          absl::StrReplaceAll(run_one_step.GetResultTokens()[0], {{"▁", " "}});
+      observer->OnNext(responses);
+    }
 
     if (ShouldStop(decode_result == kDone, benchmark_decode_token_count,
                    num_decoded_steps, executor.GetCurrentStep().value(),
@@ -533,8 +537,10 @@ absl::Status DecodeCustomSamplingStreaming(
         scores[j] += run_one_step.GetScores()[j];
       }
     }
+    if (*decode_result == kPartial) {
+      observer->OnNext(responses);
+    }
     num_decode_steps++;
-    observer->OnNext(responses);
     if (ShouldStop(*decode_result == kDone, benchmark_decode_token_count,
                    num_decode_steps, executor.GetCurrentStep().value(),
                    max_num_tokens, observer)) {

runtime/core/pipeline_test.cc

@@ -63,6 +63,7 @@ class PipelineTest : public testing::Test {
             .string());
     ASSERT_OK(tokenizer);
     tokenizer_ = std::move(*tokenizer);
+
     // The prefill tokens are the expected tokens that will be passed in at each
     // time the Prefill function is called. The values are the token ids of the
     // input prompt "Hello World!" prepended with the bos token id (2).
@@ -111,7 +112,9 @@ TEST_F(PipelineTest, Decode) {
   auto responses =
       Decode(*executor_, *tokenizer_, stop_token_detector, benchmark_info);
   EXPECT_OK(responses);
-  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?!");
+  // The response is " How's it going?" since "!" is the stop token which is
+  // not included in the response.
+  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?");
 }
 
 TEST_F(PipelineTest, DecodeReachMaxNumTokens) {
@@ -134,7 +137,9 @@ TEST_F(PipelineTest, DecodeStreaming) {
   EXPECT_OK(stop_token_detector.AddStopTokenSequence({2294}));
   EXPECT_OK(DecodeStreaming(*executor_, *tokenizer_, stop_token_detector,
                             benchmark_info, &observer));
-  EXPECT_EQ(observer.GetResponses()[0], " How's it going?!");
+  // The response is " How's it going?" since "!" is the stop token which is
+  // not included in the response.
+  EXPECT_EQ(observer.GetResponses()[0], " How's it going?");
 }
 
 TEST_F(PipelineTest, DecodeStreamingReachMaxNumTokens) {
@@ -171,18 +176,19 @@ TEST_F(PipelineTest, DecodeBytePairEncodingTokens) {
   EXPECT_CALL(*tokenizer, TokenIdsToText(std::vector<int>{18}))
       .WillOnce(testing::Return(" "));
   EXPECT_CALL(*tokenizer, TokenIdsToText(std::vector<int>{2295}))
-      .WillOnce(testing::Return("going"));
+      .WillOnce(testing::Return("going?"));
   EXPECT_CALL(*tokenizer, TokenIdsToText(std::vector<int>{2294}))
-      .WillOnce(testing::Return("?!"));
+      .WillOnce(testing::Return("!"));
 
   std::optional<BenchmarkInfo> benchmark_info;
   StopTokenDetector stop_token_detector(1);
   EXPECT_OK(stop_token_detector.AddStopTokenSequence({2294}));
   auto responses =
       Decode(*executor_, *tokenizer, stop_token_detector, benchmark_info);
   EXPECT_OK(responses);
-  // The response is truncated at the max number of tokens.
-  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?!");
+  // The response is " How's it going?" since "!" is the stop token which is
+  // not included in the response.
+  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?");
 }
 
 class PipelineCustomSamplingTest : public testing::Test {

runtime/core/session_basic_test.cc

@@ -88,7 +88,9 @@ TEST_F(SessionBasicTest, RunDecode) {
   auto responses = (*session)->RunDecode();
   EXPECT_OK(responses);
   EXPECT_EQ(responses->GetNumOutputCandidates(), 1);
-  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?!");
+  // The response is " How's it going?" since "!" is the stop token which is
+  // not included in the response.
+  EXPECT_EQ(*(responses->GetResponseTextAt(0)), " How's it going?");
 }
 
 class TestObserver : public InferenceObservable {