[to #42322933] Fix bugs for UniTE

Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11482837

modelscope/models/nlp/unite/modeling_unite.py

@@ -1,9 +1,9 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 """PyTorch UniTE model."""
 
-import math
 import warnings
 from dataclasses import dataclass
+from math import ceil
 from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
@@ -279,8 +279,8 @@ class UniTEForTranslationEvaluation(TorchModel):
         pred = self.estimator(mix_states)
         return pred.squeeze(dim=-1)
 
-    def load_checkpoint(self, path: str):
-        state_dict = torch.load(path)
+    def load_checkpoint(self, path: str, device: torch.device):
+        state_dict = torch.load(path, map_location=device)
         self.load_state_dict(state_dict)
         logger.info('Loading checkpoint parameters from %s' % path)
         return

modelscope/pipelines/nlp/translation_evaluation_pipeline.py

@@ -32,6 +32,7 @@ class TranslationEvaluationPipeline(Pipeline):
                  model: InputModel,
                  preprocessor: Optional[Preprocessor] = None,
                  eval_mode: EvaluationMode = EvaluationMode.SRC_REF,
+                 device: str = 'gpu',
                  **kwargs):
         r"""Build a translation pipeline with a model dir or a model id in the model hub.
 
@@ -51,7 +52,8 @@ class TranslationEvaluationPipeline(Pipeline):
             self.eval_mode) if preprocessor is None else preprocessor
 
         self.model.load_checkpoint(
-            osp.join(self.model.model_dir, ModelFile.TORCH_MODEL_BIN_FILE))
+            osp.join(self.model.model_dir, ModelFile.TORCH_MODEL_BIN_FILE),
+            self.device)
         self.model.eval()
 
         return

tests/pipelines/test_translation_evaluation.py

@@ -42,6 +42,16 @@ class TranslationEvaluationTest(unittest.TestCase, DemoCompatibilityCheck):
         pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
         print(pipeline_ins(input=input))
 
+        pipeline_ins = pipeline(
+            self.task, model=self.model_id_large, device='cpu')
+        print(pipeline_ins(input=input))
+
+        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
+        print(pipeline_ins(input=input))
+
+        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
+        print(pipeline_ins(input=input))
+
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_model_name_for_unite_base(self):
         input = {
@@ -68,6 +78,16 @@ class TranslationEvaluationTest(unittest.TestCase, DemoCompatibilityCheck):
         pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
         print(pipeline_ins(input=input))
 
+        pipeline_ins = pipeline(
+            self.task, model=self.model_id_base, device='cpu')
+        print(pipeline_ins(input=input))
+
+        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.SRC)
+        print(pipeline_ins(input=input))
+
+        pipeline_ins.change_eval_mode(eval_mode=EvaluationMode.REF)
+        print(pipeline_ins(input=input))
+
 
 if __name__ == '__main__':
     unittest.main()