From ec6b394d12478fb62e6fe012c86692654f0f95af Mon Sep 17 00:00:00 2001
From: liuhongwei <liuhongwei@pjlab.org.cn>
Date: Tue, 14 May 2024 20:00:59 +0800
Subject: [PATCH] Update acclerator

---
 opencompass/utils/run.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/opencompass/utils/run.py b/opencompass/utils/run.py
index 2a26b6c77..f1cb6a352 100644
--- a/opencompass/utils/run.py
+++ b/opencompass/utils/run.py
@@ -5,7 +5,8 @@
 from mmengine.config import Config
 
 from opencompass.datasets.custom import make_custom_dataset_config
-from opencompass.models import VLLM, HuggingFaceCausalLM, TurboMindModel
+from opencompass.models import (VLLM, HuggingFace, HuggingFaceCausalLM,
+                                HuggingFaceChatGLM3, TurboMindModel)
 from opencompass.partitioners import NaivePartitioner, SizePartitioner
 from opencompass.runners import DLCRunner, LocalRunner, SlurmRunner
 from opencompass.tasks import OpenICLEvalTask, OpenICLInferTask
@@ -190,7 +191,9 @@ def change_accelerator(models, accelerator):
     for model in models:
         get_logger().info(f'Transforming {model["abbr"]} to {accelerator}')
         # change HuggingFace model to VLLM or TurboMindModel
-        if model['type'] is HuggingFaceCausalLM:
+        if model['type'] in [
+                HuggingFace, HuggingFaceCausalLM, HuggingFaceChatGLM3
+        ]:
             gen_args = dict()
             if model.get('generation_kwargs') is not None:
                 generation_kwargs = model['generation_kwargs'].copy()
@@ -265,6 +268,8 @@ def change_accelerator(models, accelerator):
                         acc_model[item] = model[item]
             else:
                 raise ValueError(f'Unsupported accelerator {accelerator}')
+        else:
+            raise ValueError(f'Unsupported model type {model["type"]}')
         model_accels.append(acc_model)
     return model_accels