From 7a9018f16958ddd50071f9752d4247cdba4e87e2 Mon Sep 17 00:00:00 2001
From: wangxiyuan <wangxiyuan@huawei.com>
Date: Fri, 8 Sep 2023 11:57:01 +0800
Subject: [PATCH] Update inference.py

---
 fastchat/serve/inference.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/fastchat/serve/inference.py b/fastchat/serve/inference.py
index c97fd1c01..4e5191610 100644
--- a/fastchat/serve/inference.py
+++ b/fastchat/serve/inference.py
@@ -108,6 +108,7 @@ def generate_stream(
 
     past_key_values = out = None
     sent_interrupt = False
+    finish_reason = None
     for i in range(max_new_tokens):
         if i == 0:  # prefill
             if model.config.is_encoder_decoder:
@@ -240,12 +241,11 @@ def generate_stream(
             break
 
     # Finish stream event, which contains finish reason
-    if i == max_new_tokens - 1:
+    else:
         finish_reason = "length"
-    elif stopped:
+
+    if stopped:
         finish_reason = "stop"
-    else:
-        finish_reason = None
 
     yield {
         "text": output,