adapt to lmdeploy v0.4.0

open-compass · Apr 23, 2024 · 7344c63 · 7344c63
1 parent 81d0e4d
commit 7344c63
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/opencompass/models/lmdeploy_pytorch.py b/opencompass/models/lmdeploy_pytorch.py
@@ -145,9 +145,10 @@ def _generate(self,
         assert type(
             prompt) is str, 'We only support string for TurboMind Python API'
         input_ids = self.tokenizer.encode(prompt)
-        _, output_ids, _ = generator.infer(session_id,
+        outputs = generator.infer(session_id,
                                            input_ids,
                                            gen_config=gen_config)
+        output_ids = outputs.token_ids
         # stop engine
         if hasattr(generator, 'end'):
             generator.end(session_id)

diff --git a/opencompass/models/turbomind.py b/opencompass/models/turbomind.py
@@ -165,7 +165,7 @@ def _generate(self,
                                               sequence_end=True,
                                               step=0,
                                               stream_output=False):
-            _, output_ids, _ = outputs
+            output_ids = outputs.token_ids
             response = self.tokenizer.decode(output_ids)
             response = valid_str(response)
         # used to trim