From 561fbe9eb2184e9740a2d55ae990ff46ca4fdeee Mon Sep 17 00:00:00 2001
From: "lukeming.lkm" <lukeming.lkm@alibaba-inc.com>
Date: Wed, 2 Aug 2023 18:46:12 +0800
Subject: [PATCH] reformat stop_words_ids to sync QWen7b HF

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/13510247
---
 modelscope/models/nlp/qwen/text_generation.py | 58 ++++++++++++++++---
 1 file changed, 49 insertions(+), 9 deletions(-)

diff --git a/modelscope/models/nlp/qwen/text_generation.py b/modelscope/models/nlp/qwen/text_generation.py
index a25164cb..df8091d1 100644
--- a/modelscope/models/nlp/qwen/text_generation.py
+++ b/modelscope/models/nlp/qwen/text_generation.py
@@ -1,13 +1,15 @@
 import warnings
-from typing import Callable, Iterator, List, Optional, Tuple, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch.nn import CrossEntropyLoss
-from transformers import PreTrainedTokenizer
+from transformers import (GenerationConfig, PreTrainedTokenizer,
+                          StoppingCriteriaList)
 from transformers.generation.logits_process import LogitsProcessorList
+from transformers.generation.utils import GenerateOutput
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
 from modelscope.metainfo import Models
@@ -22,6 +24,9 @@ from .qwen_generation_utils import (BatchTokensType, HistoryType,
                                     make_context, pad_batch, switch,
                                     top_k_logits)
 
+if TYPE_CHECKING:
+    from transformers.generation.streamers import BaseStreamer
+
 logger = get_logger()
 
 
@@ -174,15 +179,9 @@ class QWenForTextGeneration(QWenPreTrainedModel):
                                             tokenizer)
         input_ids = torch.tensor([context_tokens]).to(self.device)
 
-        logits_processor_list = LogitsProcessorList([
-            StopWordsLogitsProcessor(
-                stop_words_ids=stop_words_ids,
-                eos_token_id=self.generation_config.eos_token_id)
-        ])
-
         outputs = self.generate(
             input_ids,
-            logits_processor=logits_processor_list,
+            stop_words_ids=stop_words_ids,
             return_dict_in_generate=False,
         )
 
@@ -199,3 +198,44 @@ class QWenForTextGeneration(QWenPreTrainedModel):
             history.append((query, response))
 
         return {OutputKeys.RESPONSE: response, OutputKeys.HISTORY: history}
+
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor],
+                                                    List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        streamer: Optional['BaseStreamer'] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        # Process stop_words_ids
+        stop_words_ids = kwargs.pop('stop_words_ids', None)
+        if stop_words_ids is None and generation_config is not None:
+            stop_words_ids = getattr(generation_config, 'stop_words_ids', None)
+        if stop_words_ids is None:
+            stop_words_ids = getattr(self.generation_config, 'stop_words_ids',
+                                     None)
+
+        if stop_words_ids is not None:
+            stop_words_logits_processor = StopWordsLogitsProcessor(
+                stop_words_ids=stop_words_ids,
+                eos_token_id=self.generation_config.eos_token_id)
+            if logits_processor is None:
+                logits_processor = LogitsProcessorList(
+                    [stop_words_logits_processor])
+            else:
+                logits_processor.append(stop_words_logits_processor)
+
+        return super().generate(
+            inputs,
+            generation_config,
+            logits_processor,
+            stopping_criteria,
+            prefix_allowed_tokens_fn,
+            synced_gpus,
+            streamer,
+            **kwargs,
+        )