diff --git a/modelscope/models/nlp/bert/sentence_embedding.py b/modelscope/models/nlp/bert/sentence_embedding.py index 92a9da50..b7df5ef9 100644 --- a/modelscope/models/nlp/bert/sentence_embedding.py +++ b/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 diff --git a/modelscope/models/nlp/bloom/__init__.py b/modelscope/models/nlp/bloom/__init__.py index b0f04af7..24d7202d 100644 --- a/modelscope/models/nlp/bloom/__init__.py +++ b/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( diff --git a/modelscope/models/nlp/bloom/sentence_embedding.py b/modelscope/models/nlp/bloom/sentence_embedding.py new file mode 100644 index 00000000..ec35db38 --- /dev/null +++ b/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 + `_ 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 diff --git a/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py b/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py index b03268c6..f1ca6685 100644 --- a/modelscope/preprocessors/nlp/sentence_embedding_preprocessor.py +++ b/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 diff --git a/tests/pipelines/test_sentence_embedding.py b/tests/pipelines/test_sentence_embedding.py index 13260132..a6dd89ec 100644 --- a/tests/pipelines/test_sentence_embedding.py +++ b/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)