更新sentence embedding model，支持gte，bloom sentence embedding

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

* fix linter

* bloom embedding

modelscope/models/nlp/bert/sentence_embedding.py

@@ -1,6 +1,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 
 from modelscope.metainfo import Models
@@ -61,8 +62,9 @@ class BertForSentenceEmbedding(BertPreTrainedModel):
     def __init__(self, config, **kwargs):
         super().__init__(config)
         self.config = config
-        self.pooler_type = kwargs.get('pooler_type', 'cls')
+        self.pooler_type = kwargs.get('emb_pooler_type', 'cls')
         self.pooler = Pooler(self.pooler_type)
+        self.normalize = kwargs.get('normalize', False)
         setattr(self, self.base_model_prefix,
                 BertModel(config, add_pooling_layer=False))
 
@@ -128,6 +130,8 @@ class BertForSentenceEmbedding(BertPreTrainedModel):
             output_hidden_states=output_hidden_states,
             return_dict=return_dict)
         outputs = self.pooler(outputs, attention_mask)
+        if self.normalize:
+            outputs = F.normalize(outputs, p=2, dim=-1)
         return outputs
 
     @classmethod
@@ -142,8 +146,11 @@ class BertForSentenceEmbedding(BertPreTrainedModel):
             The loaded model, which is initialized by transformers.PreTrainedModel.from_pretrained
         """
         model_dir = kwargs.get('model_dir')
-        model = super(
-            Model,
-            cls).from_pretrained(pretrained_model_name_or_path=model_dir)
+        model_kwargs = {
+            'emb_pooler_type': kwargs.get('emb_pooler_type', 'cls'),
+            'normalize': kwargs.get('normalize', False)
+        }
+        model = super(Model, cls).from_pretrained(
+            pretrained_model_name_or_path=model_dir, **model_kwargs)
         model.model_dir = model_dir
         return model

modelscope/models/nlp/bloom/__init__.py

@@ -6,10 +6,12 @@ from modelscope.utils.import_utils import LazyImportModule
 if TYPE_CHECKING:
     from .backbone import BloomModel
     from .text_generation import BloomForTextGeneration
+    from .sentence_embedding import BloomForSentenceEmbedding
 else:
     _import_structure = {
         'backbone': ['BloomModel'],
         'text_generation': ['BloomForTextGeneration'],
+        'sentence_embedding': ['BloomForSentenceEmbedding']
     }
     import sys
     sys.modules[__name__] = LazyImportModule(

modelscope/models/nlp/bloom/sentence_embedding.py

@@ -0,0 +1,165 @@
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import torch
+from transformers import BloomConfig
+from transformers import BloomModel as BloomModelTransform
+
+from modelscope.metainfo import Models
+from modelscope.models import MODELS, TorchModel
+from modelscope.outputs import SentencEmbeddingModelOutput
+from modelscope.utils.constant import Tasks
+
+
+class DecoderPooler(torch.nn.Module):
+    """
+    Parameter-free poolers to get the sentence embedding
+    'last': the last token state.
+    'weighted_mean': position weighted average of all token states.
+    """
+
+    def __init__(self, pooler_type):
+        super().__init__()
+        self.pooler_type = pooler_type
+        assert self.pooler_type in [
+            'last', 'weighted_mean'
+        ], 'unrecognized pooling type %s' % self.pooler_type
+
+    def forward(self, outputs, attention_mask):
+        last_hidden = outputs.last_hidden_state
+
+        if self.pooler_type in ['last']:
+            n, l, h = last_hidden.shape
+
+            # Get shape [n] indices of the last token (i.e. the last token for each batch item)
+            # Any sequence where min == 1, we use the entire sequence lenth since argmin = 0
+            values, indices = torch.min(attention_mask, 1, keepdim=False)
+            gather_indices = torch.where(values == 0, indices,
+                                         l) - 1  # Shape [n]
+
+            # There are empty sequences, where the index would become -1 which will crash
+            gather_indices = torch.clamp(gather_indices, min=0)
+
+            # Turn indices from shape [n] --> [n, 1, h]
+            gather_indices = gather_indices.unsqueeze(1).unsqueeze(1).expand(
+                n, 1, h)
+
+            # Gather along the 1st dim (l) (n, l, h -> n, h)
+            pooled_output = torch.gather(last_hidden, 1,
+                                         gather_indices).squeeze(dim=1)
+
+        elif self.pooler_type == 'weighted_mean':
+            input_mask_expanded = attention_mask.unsqueeze(-1).expand(
+                last_hidden.size()).float()
+            # last_hidden shape: bs, seq, hidden_dim
+            weights = (
+                torch.arange(start=1, end=last_hidden.shape[1]
+                             + 1).unsqueeze(0).unsqueeze(-1).expand(
+                                 last_hidden.size()).float().to(
+                                     last_hidden.device))
+            assert weights.shape == last_hidden.shape == input_mask_expanded.shape
+            input_mask_expanded = input_mask_expanded * weights
+
+            sum_embeddings = torch.sum(last_hidden * input_mask_expanded, 1)
+            sum_mask = input_mask_expanded.sum(1)
+            sum_mask = torch.clamp(sum_mask, min=1e-9)
+            pooled_output = sum_embeddings / sum_mask
+
+        else:
+            raise NotImplementedError
+
+        return pooled_output
+
+
+@MODELS.register_module(
+    group_key=Tasks.sentence_embedding, module_name=Models.bloom)
+class BloomForSentenceEmbedding(BloomModelTransform, TorchModel):
+    r"""
+    This model represent a text to a dense vector by the last token state or weighted mean of all token states.
+    See `Language Models are Universal Embedders
+    <https://arxiv.org/pdf/2310.08232.pdf>`_ for details.
+    """
+
+    def __init__(self, config, **kwargs):
+        super().__init__(config)
+        self.config = config
+        self.pooler_type = kwargs.get('emb_pooler_type', 'weighted_mean')
+        self.pooler = DecoderPooler(self.pooler_type)
+        self.normalize = kwargs.get('normalize', False)
+        setattr(self, self.base_model_prefix, BloomModelTransform(config))
+
+    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_udever_bloom_560m')
+            >>> preprocessor = Preprocessor.from_pretrained('damo/nlp_udever_bloom_560m')
+            >>> inputs = preprocessor({'source_sentence': ['This is a test']})
+            >>> outputs = model(**inputs)
+            >>> print(outputs)
+        """
+        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 = torch.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,
+    ):
+        outputs = self.base_model.forward(
+            input_ids, attention_mask=attention_mask)
+        embeddings = self.pooler(outputs, attention_mask)
+        if self.normalize:
+            embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=-1)
+        return embeddings
+
+    @classmethod
+    def _instantiate(cls, **kwargs):
+        """Instantiate the model.
+
+        Args:
+            kwargs: Input args.
+                    model_dir: The model dir used to load the checkpoint and the label information.
+
+        Returns:
+            The loaded model, which is initialized by transformers.PreTrainedModel.from_pretrained
+        """
+        model_dir = kwargs.get('model_dir')
+        model_kwargs = {
+            'emb_pooler_type': kwargs.get('emb_pooler_type', 'weighted_mean'),
+            'normalize': kwargs.get('normalize', False)
+        }
+        if model_dir is None:
+            config = BloomConfig(**kwargs)
+            model = cls(config)
+        else:
+            model = super(BloomModelTransform, cls).from_pretrained(
+                pretrained_model_name_or_path=model_dir, **model_kwargs)
+        model.model_dir = model_dir
+        return model

modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py

@@ -1,14 +1,19 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 
-from typing import Any, Dict
+from typing import Any, Dict, Optional
+
+import torch
 
 from modelscope.metainfo import Preprocessors
 from modelscope.preprocessors import Preprocessor
 from modelscope.preprocessors.builder import PREPROCESSORS
 from modelscope.utils.constant import Fields, ModeKeys
 from modelscope.utils.hub import get_model_type
+from modelscope.utils.logger import get_logger
 from .transformers_tokenizer import NLPTokenizer
 
+logger = get_logger()
+
 
 @PREPROCESSORS.register_module(
     Fields.nlp, module_name=Preprocessors.sentence_embedding)
@@ -46,9 +51,32 @@ class SentenceEmbeddingTransformersPreprocessor(Preprocessor):
         self.max_length = max_length
         if model_dir is not None:
             model_type = get_model_type(model_dir)
+        # we could add `boq/bod` token/prompt and `eoq/eod` token if they exist when tokenizing.
+        for k in ('boq', 'eoq', 'bod', 'eod'):
+            setattr(self, k, kwargs.pop(k, None))
         self.nlp_tokenizer = NLPTokenizer(
             model_dir, model_type, use_fast=use_fast, tokenize_kwargs=kwargs)
         super().__init__(mode=mode)
+        tokenizer = self.nlp_tokenizer.tokenizer
+        # For tokenizers like bloom
+        if tokenizer.padding_side != 'right':
+            # weighted mean pooling need pad right
+            logger.warning(
+                f'Change tokenizer.padding_side from {tokenizer.padding_side} to right'
+            )
+            tokenizer.padding_side = 'right'
+        # For decoder-only tokenizers
+        if tokenizer.pad_token is None:
+            logger.warning(
+                f'Set tokenizer.pad_token as eos_token {tokenizer.eos_token}')
+            tokenizer.pad_token = tokenizer.eos_token
+        # Currently eos is single token, we can extend to prompt later.
+        for k in ('eoq', 'eod'):
+            v = getattr(self, k, None)
+            if v is not None:
+                v = tokenizer.convert_tokens_to_ids(v)
+            setattr(self, k + '_id', v)
+        self.pad_id = tokenizer.convert_tokens_to_ids(tokenizer.pad_token)
 
     def __call__(self,
                  data: Dict,
@@ -81,13 +109,80 @@ class SentenceEmbeddingTransformersPreprocessor(Preprocessor):
         if 'return_tensors' not in kwargs:
             kwargs[
                 'return_tensors'] = 'pt' if self.mode == ModeKeys.INFERENCE else None
-        query_inputs = self.nlp_tokenizer(
-            source_sentences, padding=padding, truncation=truncation, **kwargs)
+        query_inputs = self.tokenize(
+            source_sentences,
+            is_query=True,
+            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(
+            tokenized_inputs['docs'] = self.tokenize(
                 compare_sentences,
+                is_query=kwargs.get('symmetric', False),
                 padding=padding,
                 truncation=truncation,
                 **kwargs)
         return tokenized_inputs
+
+    def tokenize(self, texts, is_query=True, return_tensors=None, **kwargs):
+        """Tokenize raw texts, add `boq/bod` token/prompt and `eoq/eod` token if they exist.
+
+        Args:
+            `texts` List[str]: texts to tokenize,
+                Example:
+                    ["how long it take to get a master's degree"]
+            `is_query` bool: whether the input text(s) is query.
+            `return_tensors` str: the `return_tensors` argument to tokenizer.
+        Returns:
+            Dict[str, Any]: the preprocessed data
+        """
+        if is_query:
+            bos, eos_id = self.boq, self.eoq_id
+        else:
+            bos, eos_id = self.bod, self.eod_id
+        if bos is not None:
+            # bos can be prompt
+            texts = [bos + t for t in texts]
+        encoding = self.nlp_tokenizer(
+            texts, return_tensors=return_tensors, **kwargs)
+        if eos_id is not None:
+            if return_tensors == 'pt':
+                self.add_eos_pt(encoding, eos_id)
+            else:
+                self.add_eos(encoding, eos_id)
+        return encoding
+
+    def add_eos_pt(self, encoding: Dict[str, torch.Tensor], eos: int):
+        """Add `eos` token id to the end of each sequence."""
+        input_ids, attn_mask = encoding['input_ids'], encoding[
+            'attention_mask']
+        batch = torch.arange(input_ids.size(0))
+        length = attn_mask.sum(-1)
+
+        if input_ids.size(1) < self.max_length:
+            ones = input_ids.new_ones(input_ids.size(0), 1)
+            attn_mask = torch.cat((ones, attn_mask), dim=1)
+            padding = ones * self.pad_id
+            input_ids = torch.cat((input_ids, padding), dim=1)
+            eos_index = length
+        else:
+            eos_index = torch.clamp(length, max=self.max_length - 1)
+            attn_mask[batch, eos_index] = 1
+        input_ids[batch, eos_index] = eos
+        encoding['input_ids'], encoding[
+            'attention_mask'] = input_ids, attn_mask
+        return
+
+    def add_eos(self, encoding: Dict[str, list], eos: int):
+        """Add `eos` token id to the end of each sequence."""
+        for ids, mask in zip(encoding['input_ids'],
+                             encoding['attention_mask']):
+            if len(mask) < self.max_length:
+                ids.append(eos)
+                mask.append(1)
+            else:
+                last = min(sum(mask), self.max_length - 1)
+                ids[last] = eos
+                mask[last] = 1
+        return

tests/pipelines/test_sentence_embedding.py

@@ -21,6 +21,7 @@ class SentenceEmbeddingTest(unittest.TestCase):
     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'
+    bloom_model_id = 'damo/udever-bloom-7b1'
 
     inputs = {
         'source_sentence': ["how long it take to get a master's degree"],
@@ -154,6 +155,14 @@ class SentenceEmbeddingTest(unittest.TestCase):
         print()
         print(f'pipeline2: {pipeline2(input=self.medical_inputs1)}')
 
+    @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
+    def test_run_with_bloom_model_from_modelhub(self):
+        model = Model.from_pretrained(self.bloom_model_id)
+        tokenizer = SentenceEmbeddingTransformersPreprocessor(model.model_dir)
+        pipeline_ins = pipeline(
+            task=Tasks.sentence_embedding, model=model, preprocessor=tokenizer)
+        print(pipeline_ins(input=self.inputs))
+
     @unittest.skipUnless(test_level() >= 0, 'skip test in current test level')
     def test_run_with_model_from_modelhub(self):
         model = Model.from_pretrained(self.model_id)