dependency version check & fix ignore_eos logic (#1099)

* version-check

* lock peft version

---------

Author: grimoire

lmdeploy/pytorch/engine/engine.py

@@ -245,6 +245,13 @@ def _on_end_session(self, reqs: Request, **kwargs):
 
     def _on_add_message(self, reqs: Request, **kwargs):
         """on add message callback."""
+
+        def __update_bad_words(msg):
+            """update bad words."""
+            sampling_param = msg.sampling_param
+            if sampling_param.ignore_eos:
+                sampling_param.bad_words.append(self.model_config.eos_token_id)
+
         for req in reqs:
             session_id = req.data['session_id']
             if session_id not in self.scheduler.sessions:
@@ -265,13 +272,15 @@ def _on_add_message(self, reqs: Request, **kwargs):
                     sampling_param=req.data['sampling_param'],
                     adapter_name=req.data['adapter_name'])
                 msg = next(iter(sess.sequences.values()))
+                __update_bad_words(msg)
                 self.scheduler.add_sequence(msg)
             else:
                 msg = next(iter(sess.sequences.values()))
                 msg.update_token_ids(req.data['token_ids'])
                 msg.remain_output_len = req.data['max_request_output_len']
                 msg.sampling_param = req.data['sampling_param']
                 msg.status = MessageStatus.WAITING
+                __update_bad_words(msg)
 
             msg.sender_id = req.sender_id
             msg.req_id = req.req_id
@@ -408,9 +417,8 @@ def _stopping_criteria(self, msg: SchedulerSequence, next_token_id: int):
         """
 
         # check eof
-        def _check_eof(sampling_param, next_token_id, eos_token_id):
-            return (not sampling_param.ignore_eos
-                    ) and next_token_id == eos_token_id
+        def _check_eof(next_token_id, eos_token_id):
+            return next_token_id == eos_token_id
 
         def _check_stop_word(sampling_param, next_token_id):
             return (sampling_param.stop_words is not None
@@ -426,8 +434,7 @@ def _check_session_len(msg, max_session_len):
             return session_len >= max_session_len
 
         sampling_param = msg.sampling_param
-        if _check_eof(sampling_param, next_token_id,
-                      self.model_config.eos_token_id):
+        if _check_eof(next_token_id, self.model_config.eos_token_id):
             return True
         if _check_stop_word(sampling_param, next_token_id):
             return True
@@ -495,8 +502,7 @@ def _can_output_token(self, token: torch.Tensor, msg: SchedulerSequence):
         """check if output is necessary."""
         if isinstance(token, torch.Tensor):
             token = token.item()
-        ignore_eos = msg.sampling_param.ignore_eos
-        if not ignore_eos and token == self.model_config.eos_token_id:
+        if token == self.model_config.eos_token_id:
             return False
 
         stop_words = msg.sampling_param.stop_words

lmdeploy/pytorch/engine/logits_process.py

@@ -61,6 +61,7 @@ def __call__(self, input_ids: torch.LongTensor,
         # bad words
         bad_words = self.sampling_param.bad_words
         if bad_words:
+            bad_words = list(set(bad_words))
             bad_words_bias = new_scores.new_zeros(new_scores.size(1))
             bad_words_bias[bad_words] = filter_value
             new_scores += bad_words_bias[None]

lmdeploy/pytorch/messages.py

@@ -35,14 +35,18 @@ def __hash__(self):
     def from_gen_config(self, gen_config: EngineGenerationConfig):
         """from gen config."""
 
+        stop_words = gen_config.stop_words or []
+        bad_words = gen_config.bad_words or []
+        if gen_config.ignore_eos:
+            bad_words += stop_words
         return SamplingParam(top_p=gen_config.top_p,
                              top_k=gen_config.top_k,
                              temperature=gen_config.temperature,
                              repetition_penalty=gen_config.repetition_penalty,
                              ignore_eos=gen_config.ignore_eos,
                              random_seed=gen_config.random_seed,
-                             stop_words=gen_config.stop_words,
-                             bad_words=gen_config.bad_words)
+                             stop_words=stop_words,
+                             bad_words=bad_words)
 
 
 class MessageStatus(enum.Enum):

requirements/runtime.txt

@@ -3,13 +3,14 @@ fire
 fuzzywuzzy
 mmengine-lite
 numpy
-peft
+peft==0.7.1
 pydantic>2.0.0
 pynvml
 safetensors
 sentencepiece
 shortuuid
 tiktoken
-torch
-transformers>=4.33.0
+torch<=2.1.2,>=2.0.0
+transformers>=4.33.0,<=4.37.1
+triton>=2.1.0,<2.2.0
 uvicorn