From 0c778c9a5ac13ebc4d561a1fe7887ea86f1aedcc Mon Sep 17 00:00:00 2001 From: "wanyu.wy" Date: Tue, 31 Jan 2023 02:20:55 +0000 Subject: [PATCH] [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 | 6 +++--- .../nlp/translation_evaluation_pipeline.py | 4 +++- .../pipelines/test_translation_evaluation.py | 20 +++++++++++++++++++ 3 files changed, 26 insertions(+), 4 deletions(-) diff --git a/modelscope/models/nlp/unite/modeling_unite.py b/modelscope/models/nlp/unite/modeling_unite.py index 6969e0c7..deea737d 100644 --- a/modelscope/models/nlp/unite/modeling_unite.py +++ b/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 diff --git a/modelscope/pipelines/nlp/translation_evaluation_pipeline.py b/modelscope/pipelines/nlp/translation_evaluation_pipeline.py index 0bd1ce9c..ba1e08a9 100644 --- a/modelscope/pipelines/nlp/translation_evaluation_pipeline.py +++ b/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 diff --git a/tests/pipelines/test_translation_evaluation.py b/tests/pipelines/test_translation_evaluation.py index 76720ac0..53524fdc 100644 --- a/tests/pipelines/test_translation_evaluation.py +++ b/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()