From e6c05a29318d781707396ae631055c4e21e4cedb Mon Sep 17 00:00:00 2001 From: "zhangyanzhao.zyz" Date: Fri, 10 Feb 2023 06:07:38 +0000 Subject: [PATCH] sentence-embedding support finetune MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit sentence-embedding模型支持finetune Link: https://code.alibaba-inc.com/Ali-MaaS/MaaS-lib/codereview/11537009 --- modelscope/metainfo.py | 5 +- .../models/nlp/bert/sentence_embedding.py | 152 ++++++++------ .../task_datasets/text_ranking_dataset.py | 2 + modelscope/outputs/nlp_outputs.py | 15 ++ .../nlp/sentence_embedding_pipeline.py | 19 +- .../nlp/sentence_embedding_preprocessor.py | 20 +- modelscope/trainers/nlp/__init__.py | 2 + .../nlp/sentence_embedding_trainer.py | 105 ++++++++++ tests/pipelines/test_sentence_embedding.py | 43 ++++ tests/pipelines/test_text_ranking.py | 16 ++ .../test_finetune_sentence_embedding.py | 187 ++++++++++++++++++ 11 files changed, 493 insertions(+), 73 deletions(-) create mode 100644 modelscope/trainers/nlp/sentence_embedding_trainer.py create mode 100644 tests/trainers/test_finetune_sentence_embedding.py diff --git a/modelscope/metainfo.py b/modelscope/metainfo.py index b5ae50d8..887dc4b7 100644 --- a/modelscope/metainfo.py +++ b/modelscope/metainfo.py @@ -455,10 +455,10 @@ DEFAULT_MODEL_FOR_PIPELINE = { Tasks.sentence_embedding: (Pipelines.sentence_embedding, 'damo/nlp_corom_sentence-embedding_english-base'), - Tasks.text_ranking: (Pipelines.text_ranking, - 'damo/nlp_corom_passage-ranking_english-base'), Tasks.text_ranking: (Pipelines.mgeo_ranking, 'damo/mgeo_address_ranking_chinese_base'), + Tasks.text_ranking: (Pipelines.text_ranking, + 'damo/nlp_corom_passage-ranking_english-base'), Tasks.word_segmentation: (Pipelines.word_segmentation, 'damo/nlp_structbert_word-segmentation_chinese-base'), @@ -765,6 +765,7 @@ class NLPTrainers(object): nlp_base_trainer = 'nlp-base-trainer' nlp_veco_trainer = 'nlp-veco-trainer' nlp_text_ranking_trainer = 'nlp-text-ranking-trainer' + nlp_sentence_embedding_trainer = 'nlp-sentence-embedding-trainer' text_generation_trainer = 'text-generation-trainer' nlp_plug_trainer = 'nlp-plug-trainer' gpt3_trainer = 'nlp-gpt3-trainer' diff --git a/modelscope/models/nlp/bert/sentence_embedding.py b/modelscope/models/nlp/bert/sentence_embedding.py index 18cecd3c..0042ee19 100644 --- a/modelscope/models/nlp/bert/sentence_embedding.py +++ b/modelscope/models/nlp/bert/sentence_embedding.py @@ -1,23 +1,111 @@ # Copyright (c) Alibaba, Inc. and its affiliates. +import torch +from torch import nn + from modelscope.metainfo import Models from modelscope.models import Model from modelscope.models.builder import MODELS -from modelscope.outputs import BackboneModelOutput +from modelscope.outputs import SentencEmbeddingModelOutput from modelscope.utils.constant import Tasks from .backbone import BertModel, BertPreTrainedModel +class Pooler(nn.Module): + """ + Parameter-free poolers to get the sentence embedding + 'cls': [CLS] representation with BERT/RoBERTa's MLP pooler. + 'cls_before_pooler': [CLS] representation without the original MLP pooler. + 'avg': average of the last layers' hidden states at each token. + 'avg_top2': average of the last two layers. + 'avg_first_last': average of the first and the last layers. + """ + + def __init__(self, pooler_type): + super().__init__() + self.pooler_type = pooler_type + assert self.pooler_type in [ + 'cls', 'avg', 'avg_top2', 'avg_first_last' + ], 'unrecognized pooling type %s' % self.pooler_type + + def forward(self, outputs, attention_mask): + last_hidden = outputs.last_hidden_state + hidden_states = outputs.hidden_states + + if self.pooler_type in ['cls']: + return last_hidden[:, 0] + elif self.pooler_type == 'avg': + return ((last_hidden * attention_mask.unsqueeze(-1)).sum(1) + / attention_mask.sum(-1).unsqueeze(-1)) + elif self.pooler_type == 'avg_first_last': + first_hidden = hidden_states[1] + last_hidden = hidden_states[-1] + pooled_result = ((first_hidden + last_hidden) / 2.0 + * attention_mask.unsqueeze(-1) + ).sum(1) / attention_mask.sum(-1).unsqueeze(-1) + return pooled_result + elif self.pooler_type == 'avg_top2': + second_last_hidden = hidden_states[-2] + last_hidden = hidden_states[-1] + pooled_result = ((last_hidden + second_last_hidden) / 2.0 + * attention_mask.unsqueeze(-1) + ).sum(1) / attention_mask.sum(-1).unsqueeze(-1) + return pooled_result + else: + raise NotImplementedError + + @MODELS.register_module(Tasks.sentence_embedding, module_name=Models.bert) class BertForSentenceEmbedding(BertPreTrainedModel): def __init__(self, config, **kwargs): super().__init__(config) self.config = config + self.pooler_type = kwargs.get('pooler_type', 'cls') + self.pooler = Pooler(self.pooler_type) setattr(self, self.base_model_prefix, BertModel(config, add_pooling_layer=False)) - def forward( + def forward(self, query=None, docs=None, labels=None): + r""" + Args: + query (:obj: `dict`): Dict of pretrained models's input for the query sequence. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + docs (:obj: `dict`): Dict of pretrained models's input for the query sequence. See + :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for + details. + Returns: + Returns `modelscope.outputs.SentencEmbeddingModelOutput + Examples: + >>> from modelscope.models import Model + >>> from modelscope.preprocessors import Preprocessor + >>> model = Model.from_pretrained('damo/nlp_corom_sentence-embedding_chinese-base') + >>> preprocessor = Preprocessor.from_pretrained('damo/nlp_corom_sentence-embedding_chinese-base') + >>> print(model(**preprocessor('source_sentence':['This is a test']))) + """ + query_embeddings, doc_embeddings = None, None + if query is not None: + query_embeddings = self.encode(**query) + if docs is not None: + doc_embeddings = self.encode(**docs) + outputs = SentencEmbeddingModelOutput( + query_embeddings=query_embeddings, doc_embeddings=doc_embeddings) + if query_embeddings is None or doc_embeddings is None: + return outputs + if self.base_model.training: + loss_fct = nn.CrossEntropyLoss() + scores = torch.matmul(query_embeddings, doc_embeddings.T) + if labels is None: + labels = torch.arange( + scores.size(0), device=scores.device, dtype=torch.long) + labels = labels * ( + doc_embeddings.size(0) // query_embeddings.size(0)) + loss = loss_fct(scores, labels) + outputs.loss = loss + return outputs + + def encode( self, input_ids=None, attention_mask=None, @@ -28,62 +116,8 @@ class BertForSentenceEmbedding(BertPreTrainedModel): output_attentions=None, output_hidden_states=None, return_dict=None, - ) -> BackboneModelOutput: - r""" - Args: - input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`): - Indices of input sequence tokens in the vocabulary. - - Indices can be obtained using :class:`~modelscope.models.nlp.structbert.SbertTokenizer`. See - :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for - details. - - attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): - Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``: - - - 1 for tokens that are **not masked**, - - 0 for tokens that are **masked**. - - token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): - Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0, - 1]``: - - - 0 corresponds to a `sentence A` token, - - 1 corresponds to a `sentence B` token. - - position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`): - Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0, - config.max_position_embeddings - 1]``. - - head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`): - Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``: - - - 1 indicates the head is **not masked**, - - 0 indicates the head is **masked**. - - inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`): - Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation. - This is useful if you want more control over how to convert :obj:`input_ids` indices into associated - vectors than the model's internal embedding lookup matrix. - output_attentions (:obj:`bool`, `optional`): - Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned - tensors for more detail. - output_hidden_states (:obj:`bool`, `optional`): - Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for - more detail. - return_dict (:obj:`bool`, `optional`): - Whether or not to return a :class:`~transformers.ModelOutput` instead of a plain tuple. - Returns: - Returns `modelscope.outputs.AttentionTextClassificationModelOutput` - - Examples: - >>> from modelscope.models import Model - >>> from modelscope.preprocessors import Preprocessor - >>> model = Model.from_pretrained('damo/nlp_corom_sentence-embedding_chinese-base') - >>> preprocessor = Preprocessor.from_pretrained('damo/nlp_corom_sentence-embedding_chinese-base') - >>> print(model(**preprocessor('This is a test'))) - """ - return self.base_model.forward( + ): + outputs = self.base_model.forward( input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, @@ -93,6 +127,8 @@ class BertForSentenceEmbedding(BertPreTrainedModel): output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=return_dict) + outputs = self.pooler(outputs, attention_mask) + return outputs @classmethod def _instantiate(cls, **kwargs): diff --git a/modelscope/msdatasets/task_datasets/text_ranking_dataset.py b/modelscope/msdatasets/task_datasets/text_ranking_dataset.py index 54276843..19f07110 100644 --- a/modelscope/msdatasets/task_datasets/text_ranking_dataset.py +++ b/modelscope/msdatasets/task_datasets/text_ranking_dataset.py @@ -17,6 +17,8 @@ from .torch_base_dataset import TorchTaskDataset @TASK_DATASETS.register_module( group_key=Tasks.text_ranking, module_name=Models.bert) +@TASK_DATASETS.register_module( + group_key=Tasks.sentence_embedding, module_name=Models.bert) class TextRankingDataset(TorchTaskDataset): def __init__(self, diff --git a/modelscope/outputs/nlp_outputs.py b/modelscope/outputs/nlp_outputs.py index e285d40f..7bada5dc 100644 --- a/modelscope/outputs/nlp_outputs.py +++ b/modelscope/outputs/nlp_outputs.py @@ -415,3 +415,18 @@ class DialogueUserSatisfactionEstimationModelOutput(ModelOutputBase): logits (`Tensor`): The logits output of the model. """ logits: Tensor = None + + +@dataclass +class SentencEmbeddingModelOutput(ModelOutputBase): + """The output class for text classification models. + + Args: + query_embs (`Tensor`, *optional*): The tensor of the query embeddings. + doc_embs (`Tensor`, *optional*) Then tensor of the doc embeddings. + loss (`torch.FloatTensor` of shape `(1,)`, *optional*): Sentence Embedding modeling loss. + """ + + query_embeddings: Tensor = None + doc_embeddings: Tensor = None + loss: Tensor = None diff --git a/modelscope/pipelines/nlp/sentence_embedding_pipeline.py b/modelscope/pipelines/nlp/sentence_embedding_pipeline.py index e8132868..7aaa073b 100644 --- a/modelscope/pipelines/nlp/sentence_embedding_pipeline.py +++ b/modelscope/pipelines/nlp/sentence_embedding_pipeline.py @@ -3,6 +3,7 @@ from typing import Any, Dict, Optional, Union import numpy as np +import torch from modelscope.metainfo import Pipelines from modelscope.models import Model @@ -65,11 +66,17 @@ class SentenceEmbeddingPipeline(Pipeline): Returns: Dict[str, Any]: the predicted text representation """ - embs = inputs['last_hidden_state'][:, 0].cpu().numpy() - num_sent = embs.shape[0] - if num_sent >= 2: - scores = np.dot(embs[0:1, ], np.transpose(embs[1:, ], - (1, 0))).tolist()[0] + embeddings = inputs['query_embeddings'] + doc_embeddings = inputs['doc_embeddings'] + if doc_embeddings is not None: + embeddings = torch.cat((embeddings, doc_embeddings), dim=0) + embeddings = embeddings.detach().cpu().numpy() + if doc_embeddings is not None: + scores = np.dot(embeddings[0:1, ], + np.transpose(embeddings[1:, ], (1, 0))).tolist()[0] else: scores = [] - return {OutputKeys.TEXT_EMBEDDING: embs, OutputKeys.SCORES: scores} + return { + OutputKeys.TEXT_EMBEDDING: embeddings, + OutputKeys.SCORES: scores + } diff --git a/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py b/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py index 5930e007..b03268c6 100644 --- a/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py +++ b/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py @@ -43,6 +43,7 @@ class SentenceEmbeddingTransformersPreprocessor(Preprocessor): 'sequence_length', 128) kwargs.pop('sequence_length', None) model_type = None + self.max_length = max_length if model_dir is not None: model_type = get_model_type(model_dir) self.nlp_tokenizer = NLPTokenizer( @@ -72,16 +73,21 @@ class SentenceEmbeddingTransformersPreprocessor(Preprocessor): """ source_sentences = data[self.first_sequence] if self.second_sequence in data: + if isinstance(source_sentences[0], list): + source_sentences = [source_sentences[0]] compare_sentences = data[self.second_sequence] - sentences = [source_sentences[0]] - for sent in compare_sentences: - sentences.append(sent) else: - sentences = source_sentences + compare_sentences = None if 'return_tensors' not in kwargs: kwargs[ 'return_tensors'] = 'pt' if self.mode == ModeKeys.INFERENCE else None - - tokenized_inputs = self.nlp_tokenizer( - sentences, padding=padding, truncation=truncation, **kwargs) + query_inputs = self.nlp_tokenizer( + source_sentences, padding=padding, truncation=truncation, **kwargs) + tokenized_inputs = {'query': query_inputs, 'docs': None} + if compare_sentences is not None and len(compare_sentences) > 0: + tokenized_inputs['docs'] = self.nlp_tokenizer( + compare_sentences, + padding=padding, + truncation=truncation, + **kwargs) return tokenized_inputs diff --git a/modelscope/trainers/nlp/__init__.py b/modelscope/trainers/nlp/__init__.py index e3c39cf2..125e82c6 100644 --- a/modelscope/trainers/nlp/__init__.py +++ b/modelscope/trainers/nlp/__init__.py @@ -8,12 +8,14 @@ if TYPE_CHECKING: from .csanmt_translation_trainer import CsanmtTranslationTrainer from .text_ranking_trainer import TextRankingTrainer from .text_generation_trainer import TextGenerationTrainer + from .sentence_embedding_trainer import SentenceEmbeddingTrainer else: _import_structure = { 'sequence_classification_trainer': ['SequenceClassificationTrainer'], 'csanmt_translation_trainer': ['CsanmtTranslationTrainer'], 'text_ranking_trainer': ['TextRankingTrainer'], 'text_generation_trainer': ['TextGenerationTrainer'], + 'sentence_emebedding_trainer': ['SentenceEmbeddingTrainer'] } import sys diff --git a/modelscope/trainers/nlp/sentence_embedding_trainer.py b/modelscope/trainers/nlp/sentence_embedding_trainer.py new file mode 100644 index 00000000..b2116443 --- /dev/null +++ b/modelscope/trainers/nlp/sentence_embedding_trainer.py @@ -0,0 +1,105 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. + +import time +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple, Union + +import numpy as np +import torch +from torch import nn +from torch.utils.data import DataLoader, Dataset +from tqdm import tqdm +from transformers import DataCollatorWithPadding + +from modelscope.metainfo import Trainers +from modelscope.models.base import Model, TorchModel +from modelscope.models.nlp import BertForTextRanking +from modelscope.msdatasets.ms_dataset import MsDataset +from modelscope.preprocessors.base import Preprocessor +from modelscope.trainers.builder import TRAINERS +from modelscope.trainers.nlp_trainer import NlpEpochBasedTrainer +from modelscope.utils.constant import DEFAULT_MODEL_REVISION +from modelscope.utils.logger import get_logger + +logger = get_logger() + + +@dataclass +class SentenceEmbeddingCollator(DataCollatorWithPadding): + """ + Wrapper that does conversion from List[Tuple[encode_qry, encode_psg]] to List[qry], List[psg] + and pass batch separately to the actual collator. + Abstract out data detail for the model. + """ + max_length = 128 + tokenizer = None + + def __call__(self, features): + qq = [f['query'] for f in features] + dd = [f['docs'] for f in features] + keys = qq[0].keys() + qq = {k: [ele[k] for ele in qq] for k in keys} + q_collated = self.tokenizer._tokenizer.pad( + qq, + padding='max_length', + max_length=self.max_length, + return_tensors='pt') + keys = dd[0].keys() + dd = {k: sum([ele[k] for ele in dd], []) for k in keys} + d_collated = self.tokenizer._tokenizer.pad( + dd, + padding='max_length', + max_length=self.max_length, + return_tensors='pt') + return {'query': q_collated, 'docs': d_collated} + + +@TRAINERS.register_module(module_name=Trainers.nlp_sentence_embedding_trainer) +class SentenceEmbeddingTrainer(NlpEpochBasedTrainer): + + def __init__( + self, + model: Optional[Union[TorchModel, nn.Module, str]] = None, + cfg_file: Optional[str] = None, + cfg_modify_fn: Optional[Callable] = None, + arg_parse_fn: Optional[Callable] = None, + data_collator: Optional[Callable] = None, + train_dataset: Optional[Union[MsDataset, Dataset]] = None, + eval_dataset: Optional[Union[MsDataset, Dataset]] = None, + preprocessor: Optional[Preprocessor] = None, + optimizers: Tuple[torch.optim.Optimizer, + torch.optim.lr_scheduler._LRScheduler] = (None, + None), + model_revision: Optional[str] = DEFAULT_MODEL_REVISION, + **kwargs): + + super().__init__( + model=model, + cfg_file=cfg_file, + cfg_modify_fn=cfg_modify_fn, + arg_parse_fn=arg_parse_fn, + data_collator=data_collator, + preprocessor=preprocessor, + optimizers=optimizers, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + model_revision=model_revision, + **kwargs) + + def get_data_collator(self, data_collator): + """Get the data collator for both training and evaluating. + + Args: + data_collator: The input data_collator param. + + Returns: + The train_data_collator and eval_data_collator, can be None. + """ + if data_collator is None: + data_collator = SentenceEmbeddingCollator( + tokenizer=self.train_preprocessor.nlp_tokenizer, + max_length=self.train_preprocessor.max_length) + return super().get_data_collator(data_collator) + + def evauate(self): + return {} diff --git a/tests/pipelines/test_sentence_embedding.py b/tests/pipelines/test_sentence_embedding.py index 06be4850..13260132 100644 --- a/tests/pipelines/test_sentence_embedding.py +++ b/tests/pipelines/test_sentence_embedding.py @@ -14,9 +14,13 @@ from modelscope.utils.test_utils import test_level class SentenceEmbeddingTest(unittest.TestCase): model_id = 'damo/nlp_corom_sentence-embedding_english-base' + tiny_model_id = 'damo/nlp_corom_sentence-embedding_english-tiny' ecom_base_model_id = 'damo/nlp_corom_sentence-embedding_chinese-base-ecom' + ecom_tiny_model_id = 'damo/nlp_corom_sentence-embedding_chinese-tiny-ecom' medical_base_model_id = 'damo/nlp_corom_sentence-embedding_chinese-base-medical' + medical_tiny_model_id = 'damo/nlp_corom_sentence-embedding_chinese-tiny-medical' general_base_model_id = 'damo/nlp_corom_sentence-embedding_chinese-base' + general_tiny_model_id = 'damo/nlp_corom_sentence-embedding_chinese-tiny' inputs = { 'source_sentence': ["how long it take to get a master's degree"], @@ -51,6 +55,25 @@ class SentenceEmbeddingTest(unittest.TestCase): ] } + general_inputs1 = { + 'source_sentence': ['功和功率的区别'], + 'sentences_to_compare': [ + '功反映做功多少，功率反映做功快慢。', + '什么是有功功率和无功功率?无功功率有什么用什么是有功功率和无功功率?无功功率有什么用电力系统中的电源是由发电机产生的三相正弦交流电,在交>流电路中,由电源供给负载的电功率有两种;一种是有功功率,一种是无功功率。', + '优质解答在物理学中,用电功率表示消耗电能的快慢．电功率用P表示,它的单位是瓦特（Watt）,简称瓦（Wa）符号是W.电流在单位时间内做的功叫做电功率 以灯泡为例,电功率越大\ + ,灯泡越亮.灯泡的亮暗由电功率（实际功率）决定,不由通过的电流、电压、电能决定!', + ] + } + + general_inputs2 = { + 'source_sentence': [ + '功反映做功多少，功率反映做功快慢。', + '什么是有功功率和无功功率?无功功率有什么用什么是有功功率和无功功率?无功功率有什么用电力系统中的电源是由发电机产生的三相正弦交流电,在交>流电路中,由电源供给负载的电功率有两种;一种是有功功率,一种是无功功率。', + '优质解答在物理学中,用电功率表示消耗电能的快慢．电功率用P表示,它的单位是瓦特（Watt）,简称瓦（Wa）符号是W.电流在单位时间内做的功叫做电功率 以灯泡为例,电功率越大\ + ,灯泡越亮.灯泡的亮暗由电功率（实际功率）决定,不由通过的电流、电压、电能决定!', + ] + } + ecom_inputs1 = { 'source_sentence': ['毛绒玩具'], 'sentences_to_compare': ['大熊泰迪熊猫毛绒玩具公仔布娃娃抱抱熊', '背心式狗狗牵引绳'] @@ -144,6 +167,9 @@ class SentenceEmbeddingTest(unittest.TestCase): pipeline_ins = pipeline( task=Tasks.sentence_embedding, model=self.model_id) print(pipeline_ins(input=self.inputs)) + pipeline_ins = pipeline( + task=Tasks.sentence_embedding, model=self.tiny_model_id) + print(pipeline_ins(input=self.inputs)) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_with_default_model(self): @@ -155,12 +181,18 @@ class SentenceEmbeddingTest(unittest.TestCase): pipeline_ins = pipeline( task=Tasks.sentence_embedding, model=self.ecom_base_model_id) print(pipeline_ins(input=self.ecom_inputs2)) + pipeline_ins = pipeline( + task=Tasks.sentence_embedding, model=self.ecom_tiny_model_id) + print(pipeline_ins(input=self.ecom_inputs2)) @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_medical_model_with_model_name(self): pipeline_ins = pipeline( task=Tasks.sentence_embedding, model=self.medical_base_model_id) print(pipeline_ins(input=self.medical_inputs1)) + pipeline_ins = pipeline( + task=Tasks.sentence_embedding, model=self.medical_tiny_model_id) + print(pipeline_ins(input=self.medical_inputs1)) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_with_el_model(self): @@ -168,6 +200,17 @@ class SentenceEmbeddingTest(unittest.TestCase): task=Tasks.sentence_embedding, model=self.el_model_id) print(pipeline_ins(input=self.el_inputs)) + @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') + def test_run_general_model_with_model_name(self): + pipeline_ins = pipeline( + task=Tasks.sentence_embedding, model=self.general_base_model_id) + print(pipeline_ins(input=self.general_inputs1)) + print(pipeline_ins(input=self.general_inputs2)) + pipeline_ins = pipeline( + task=Tasks.sentence_embedding, model=self.general_tiny_model_id) + print(pipeline_ins(input=self.general_inputs1)) + print(pipeline_ins(input=self.general_inputs2)) + if __name__ == '__main__': unittest.main() diff --git a/tests/pipelines/test_text_ranking.py b/tests/pipelines/test_text_ranking.py index 40c4da58..26197410 100644 --- a/tests/pipelines/test_text_ranking.py +++ b/tests/pipelines/test_text_ranking.py @@ -14,6 +14,7 @@ from modelscope.utils.test_utils import test_level class TextRankingTest(unittest.TestCase): base_model_id = 'damo/nlp_corom_passage-ranking_english-base' + tiny_model_id = 'damo/nlp_corom_passage-ranking_english-tiny' inputs = { 'source_sentence': ["how long it take to get a master's degree"], 'sentences_to_compare': [ @@ -25,6 +26,7 @@ class TextRankingTest(unittest.TestCase): } chinese_base_model_id = 'damo/nlp_rom_passage-ranking_chinese-base' + chinese_tiny_model_id = 'damo/nlp_corom_passage-ranking_chinese-tiny' chinese_inputs = { 'source_sentence': ['功和功率的区别'], 'sentences_to_compare': [ @@ -36,12 +38,14 @@ class TextRankingTest(unittest.TestCase): } ecom_base_model_id = 'damo/nlp_corom_passage-ranking_chinese-base-ecom' + ecom_tiny_model_id = 'damo/nlp_corom_passage-ranking_chinese-tiny-ecom' ecom_inputs = { 'source_sentence': ['毛绒玩具'], 'sentences_to_compare': ['大熊泰迪熊猫毛绒玩具公仔布娃娃抱抱熊', '背心式狗狗牵引绳'] } medical_base_model_id = 'damo/nlp_corom_passage-ranking_chinese-base-medical' + medical_tiny_model_id = 'damo/nlp_corom_passage-ranking_chinese-tiny-medical' medical_inputs = { 'source_sentence': ['肠道不适可以服用益生菌吗'], 'sentences_to_compare': ['肠胃不好能吃益生菌,益生菌有调节肠胃道菌群的作用', '身体发烧应该多喝水'] @@ -86,6 +90,9 @@ class TextRankingTest(unittest.TestCase): pipeline_ins = pipeline( task=Tasks.text_ranking, model=self.base_model_id) print(pipeline_ins(input=self.inputs)) + pipeline_ins = pipeline( + task=Tasks.text_ranking, model=self.tiny_model_id) + print(pipeline_ins(input=self.inputs)) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_with_default_model(self): @@ -97,18 +104,27 @@ class TextRankingTest(unittest.TestCase): pipeline_ins = pipeline( task=Tasks.text_ranking, model=self.chinese_base_model_id) print(pipeline_ins(input=self.chinese_inputs)) + pipeline_ins = pipeline( + task=Tasks.text_ranking, model=self.chinese_tiny_model_id) + print(pipeline_ins(input=self.chinese_inputs)) @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_ecom_model_with_model_name(self): pipeline_ins = pipeline( task=Tasks.text_ranking, model=self.ecom_base_model_id) print(pipeline_ins(input=self.ecom_inputs)) + pipeline_tiny_ins = pipeline( + task=Tasks.text_ranking, model=self.ecom_tiny_model_id) + print(pipeline_tiny_ins(input=self.ecom_inputs)) @unittest.skipUnless(test_level() >= 0, 'skip test in current test level') def test_run_medical_model_with_model_name(self): pipeline_ins = pipeline( task=Tasks.text_ranking, model=self.medical_base_model_id) print(pipeline_ins(input=self.medical_inputs)) + pipeline_ins = pipeline( + task=Tasks.text_ranking, model=self.medical_tiny_model_id) + print(pipeline_ins(input=self.medical_inputs)) @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') def test_run_with_el_model(self): diff --git a/tests/trainers/test_finetune_sentence_embedding.py b/tests/trainers/test_finetune_sentence_embedding.py new file mode 100644 index 00000000..f0ca80cd --- /dev/null +++ b/tests/trainers/test_finetune_sentence_embedding.py @@ -0,0 +1,187 @@ +# Copyright (c) Alibaba, Inc. and its affiliates. +import os +import shutil +import tempfile +import unittest +from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union + +import torch +from transformers.tokenization_utils_base import PreTrainedTokenizerBase + +from modelscope.metainfo import Trainers +from modelscope.models import Model +from modelscope.msdatasets import MsDataset +from modelscope.pipelines import pipeline +from modelscope.trainers import build_trainer +from modelscope.utils.constant import ModelFile, Tasks +from modelscope.utils.test_utils import test_level + + +class TestFinetuneSentenceEmbedding(unittest.TestCase): + inputs = { + 'source_sentence': ["how long it take to get a master's degree"], + 'sentences_to_compare': [ + "On average, students take about 18 to 24 months to complete a master's degree.", + 'On the other hand, some students prefer to go at a slower pace and choose to take ' + 'several years to complete their studies.', + 'It can take anywhere from two semesters' + ] + } + + def setUp(self): + print(('Testing %s.%s' % (type(self).__name__, self._testMethodName))) + self.tmp_dir = tempfile.TemporaryDirectory().name + if not os.path.exists(self.tmp_dir): + os.makedirs(self.tmp_dir) + + def tearDown(self): + shutil.rmtree(self.tmp_dir) + super().tearDown() + + def finetune(self, + model_id, + train_dataset, + eval_dataset, + name=Trainers.nlp_sentence_embedding_trainer, + cfg_modify_fn=None, + **kwargs): + kwargs = dict( + model=model_id, + train_dataset=train_dataset, + eval_dataset=eval_dataset, + work_dir=self.tmp_dir, + cfg_modify_fn=cfg_modify_fn, + **kwargs) + + os.environ['LOCAL_RANK'] = '0' + trainer = build_trainer(name=name, default_args=kwargs) + trainer.train() + + @unittest.skipUnless(test_level() >= 1, 'skip test in current test level') + def test_finetune_msmarco(self): + + def cfg_modify_fn(cfg): + neg_sample = 2 + cfg.task = 'sentence-embedding' + cfg['preprocessor'] = {'type': 'sentence-embedding'} + cfg.train.optimizer.lr = 2e-5 + cfg['dataset'] = { + 'train': { + 'type': 'bert', + 'query_sequence': 'query', + 'pos_sequence': 'positive_passages', + 'neg_sequence': 'negative_passages', + 'text_fileds': ['title', 'text'], + 'qid_field': 'query_id', + 'neg_sample': neg_sample + }, + 'val': { + 'type': 'bert', + 'query_sequence': 'query', + 'pos_sequence': 'positive_passages', + 'neg_sequence': 'negative_passages', + 'text_fileds': ['title', 'text'], + 'qid_field': 'query_id' + }, + } + cfg['evaluation']['dataloader']['batch_size_per_gpu'] = 30 + cfg.train.max_epochs = 1 + cfg.train.train_batch_size = 2 + cfg.train.lr_scheduler = { + 'type': 'LinearLR', + 'start_factor': 1.0, + 'end_factor': 0.0, + 'options': { + 'by_epoch': False + } + } + cfg.model['neg_sample'] = 4 + cfg.train.hooks = [{ + 'type': 'CheckpointHook', + 'interval': 1 + }, { + 'type': 'TextLoggerHook', + 'interval': 1 + }, { + 'type': 'IterTimerHook' + }] + return cfg + + # load dataset + ds = MsDataset.load('passage-ranking-demo', 'zyznull') + train_ds = ds['train'].to_hf_dataset() + dev_ds = ds['dev'].to_hf_dataset() + + model_id = 'damo/nlp_corom_sentence-embedding_english-base' + self.finetune( + model_id=model_id, + train_dataset=train_ds, + eval_dataset=dev_ds, + cfg_modify_fn=cfg_modify_fn) + + output_dir = os.path.join(self.tmp_dir, ModelFile.TRAIN_OUTPUT_DIR) + self.pipeline_sentence_embedding(output_dir) + + @unittest.skipUnless(test_level() >= 2, 'skip test in current test level') + def test_finetune_dureader(self): + + def cfg_modify_fn(cfg): + cfg.task = 'sentence-embedding' + cfg['preprocessor'] = { + 'type': 'sentence-embedding', + 'max_length': 384 + } + cfg.train.optimizer.lr = 3e-5 + cfg['dataset'] = { + 'train': { + 'type': 'bert', + 'query_sequence': 'query', + 'pos_sequence': 'positive_passages', + 'neg_sequence': 'negative_passages', + 'text_fileds': ['text'], + 'qid_field': 'query_id', + 'neg_sample': 4 + }, + 'val': { + 'type': 'bert', + 'query_sequence': 'query', + 'pos_sequence': 'positive_passages', + 'neg_sequence': 'negative_passages', + 'text_fileds': ['text'], + 'qid_field': 'query_id' + }, + } + cfg['evaluation']['dataloader']['batch_size_per_gpu'] = 3 + cfg.train.max_epochs = 2 + cfg.train.train_batch_size = 4 + cfg.train.hooks = [{ + 'type': 'CheckpointHook', + 'interval': 1 + }, { + 'type': 'TextLoggerHook', + 'interval': 1 + }, { + 'type': 'IterTimerHook' + }] + return cfg + + # load dataset + ds = MsDataset.load('dureader-retrieval-ranking', 'zyznull') + train_ds = ds['train'].to_hf_dataset().shard(1000, index=0) + dev_ds = ds['dev'].to_hf_dataset() + model_id = 'damo/nlp_corom_sentence-embedding_chinese-base' + self.finetune( + model_id=model_id, + train_dataset=train_ds, + eval_dataset=dev_ds, + cfg_modify_fn=cfg_modify_fn) + + def pipeline_sentence_embedding(self, model_dir): + model = Model.from_pretrained(model_dir) + pipeline_ins = pipeline(task=Tasks.sentence_embedding, model=model) + print('inputs', self.inputs) + print(pipeline_ins(input=self.inputs)) + + +if __name__ == '__main__': + unittest.main()