sentence-embedding support finetune

sentence-embedding模型支持finetune

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

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'

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):

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,

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

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
+        }

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

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

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 {}

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()

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):

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()