[Fix] Fix turbomind_tis (#992)

open-compass · Mar 22, 2024 · e415ddf · e415ddf
1 parent 054e9fa
commit e415ddf
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 1 deletion.
diff --git a/configs/eval_internlm_chat_turbomind_tis.py b/configs/eval_internlm_chat_turbomind_tis.py
@@ -6,7 +6,7 @@
     from .datasets.mmlu.mmlu_gen_a484b3 import mmlu_datasets
     from .datasets.ceval.ceval_gen_5f30c7 import ceval_datasets
     from .datasets.SuperGLUE_WiC.SuperGLUE_WiC_gen_d06864 import WiC_datasets
-    from .datasets.SuperGLUE_WSC.SuperGLUE_WSC_gen_6dc406 import WSC_datasets
+    from .datasets.SuperGLUE_WSC.SuperGLUE_WSC_gen_7902a7 import WSC_datasets
     from .datasets.triviaqa.triviaqa_gen_2121ce import triviaqa_datasets
     from .datasets.gsm8k.gsm8k_gen_1d7fe4 import gsm8k_datasets
     from .datasets.humaneval.humaneval_gen_8e312c import humaneval_datasets

diff --git a/opencompass/models/turbomind_tis.py b/opencompass/models/turbomind_tis.py
@@ -47,6 +47,8 @@ def __init__(
         super().__init__(path=path,
                          max_seq_len=max_seq_len,
                          meta_template=meta_template)
+        from lmdeploy.serve.turbomind.utils import Preprocessor
+        self.preprocess = Preprocessor(tis_addr)
         self.logger = get_logger()
         self.template_parser = LMTemplateParser(meta_template)
         self.eos_token_id = None
@@ -83,6 +85,10 @@ def generate(
                              [temperature] * len(inputs)))
         return results
 
+    def get_token_len(self, prompt: str) -> int:
+        input_ids, _ = self.preprocess(prompt)
+        return input_ids.shape[-1]
+
     def wait(self):
         """Wait till the next query can be sent.