transformers
copied
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Readme
Files and versions
260 lines
9.5 KiB
260 lines
9.5 KiB
# Copyright 2021 Zilliz. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import numpy
|
|
import os
|
|
import requests
|
|
import torch
|
|
import shutil
|
|
from pathlib import Path
|
|
from typing import Union
|
|
from collections import OrderedDict
|
|
|
|
from transformers import AutoTokenizer, AutoConfig, AutoModel
|
|
|
|
from towhee.operator import NNOperator
|
|
from towhee import register
|
|
# from towhee.dc2 import accelerate
|
|
|
|
import warnings
|
|
import logging
|
|
from transformers import logging as t_logging
|
|
|
|
log = logging.getLogger('run_op')
|
|
warnings.filterwarnings('ignore')
|
|
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
|
|
t_logging.set_verbosity_error()
|
|
|
|
|
|
# @accelerate
|
|
class Model:
|
|
def __init__(self, model):
|
|
self.model = model
|
|
|
|
def __call__(self, *args, **kwargs):
|
|
outs = self.model(*args, **kwargs, return_dict=True)
|
|
return outs['last_hidden_state']
|
|
|
|
|
|
@register(output_schema=['vec'])
|
|
class AutoTransformers(NNOperator):
|
|
"""
|
|
NLP embedding operator that uses the pretrained transformers model gathered by huggingface.
|
|
Args:
|
|
model_name (`str`):
|
|
The model name to load a pretrained model from transformers.
|
|
checkpoint_path (`str`):
|
|
The local checkpoint path.
|
|
tokenizer (`object`):
|
|
The tokenizer to tokenize input text as model inputs.
|
|
"""
|
|
|
|
def __init__(self,
|
|
model_name: str = None,
|
|
checkpoint_path: str = None,
|
|
tokenizer: object = None,
|
|
device: str = None,
|
|
norm: bool = False
|
|
):
|
|
super().__init__()
|
|
if device:
|
|
self.device = device
|
|
else:
|
|
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
self.model_name = self.map_model_names(model_name)
|
|
print(1111, self.model_name)
|
|
if tokenizer:
|
|
self.tokenizer = tokenizer
|
|
else:
|
|
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
|
if not self.tokenizer.pad_token:
|
|
self.tokenizer.pad_token = '[PAD]'
|
|
self.norm = norm
|
|
self.checkpoint_path = checkpoint_path
|
|
|
|
if self.model_name:
|
|
model_list = self.supported_model_names()
|
|
# assert model_name in model_list, f"Invalid model name: {model_name}. Supported model names: {model_list}"
|
|
self.model = Model(self._model)
|
|
else:
|
|
log.warning('The operator is initialized without specified model.')
|
|
pass
|
|
|
|
def __call__(self, data: Union[str, list]) -> numpy.ndarray:
|
|
if isinstance(data, str):
|
|
txt = [data]
|
|
else:
|
|
txt = data
|
|
try:
|
|
inputs = self.tokenizer(txt, padding=True, truncation=True, return_tensors='pt').to(self.device)
|
|
except Exception as e:
|
|
log.error(f'Fail to tokenize inputs: {e}')
|
|
raise e
|
|
try:
|
|
outs = self.model(**inputs)
|
|
except Exception as e:
|
|
log.error(f'Invalid input for the model: {self.model_name}')
|
|
raise e
|
|
outs = self.post_proc(outs, inputs)
|
|
if self.norm:
|
|
outs = torch.nn.functional.normalize(outs, )
|
|
features = outs.cpu().detach().numpy()
|
|
if isinstance(data, str):
|
|
features = features.squeeze(0)
|
|
else:
|
|
features = list(features)
|
|
return features
|
|
|
|
@property
|
|
def _model(self):
|
|
model = AutoModel.from_pretrained(self.model_name).to(self.device)
|
|
if hasattr(model, 'pooler') and model.pooler:
|
|
model.pooler = None
|
|
if self.checkpoint_path:
|
|
try:
|
|
state_dict = torch.load(self.checkpoint_path, map_location=self.device)
|
|
model.load_state_dict(state_dict)
|
|
except Exception:
|
|
log.error(f'Fail to load weights from {self.checkpoint_path}')
|
|
model.eval()
|
|
return model
|
|
|
|
@property
|
|
def model_config(self):
|
|
configs = AutoConfig.from_pretrained(self.model_name)
|
|
return configs
|
|
|
|
@property
|
|
def onnx_config(self):
|
|
from transformers.onnx.features import FeaturesManager
|
|
model_kind, model_onnx_config = FeaturesManager.check_supported_model_or_raise(
|
|
self._model, feature='default')
|
|
old_config = model_onnx_config(self.model_config)
|
|
onnx_config = {
|
|
'inputs': dict(old_config.inputs),
|
|
'outputs': {'last_hidden_state': old_config.outputs['last_hidden_state']}
|
|
}
|
|
return onnx_config
|
|
|
|
def post_proc(self, token_embeddings, inputs):
|
|
token_embeddings = token_embeddings.to(self.device)
|
|
attention_mask = inputs['attention_mask'].to(self.device)
|
|
input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
|
|
sentence_embs = torch.sum(
|
|
token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
|
|
return sentence_embs
|
|
|
|
def save_model(self, model_type: str = 'pytorch', output_file: str = 'default'):
|
|
if output_file == 'default':
|
|
output_file = str(Path(__file__).parent)
|
|
output_file = os.path.join(output_file, 'saved', model_type)
|
|
os.makedirs(output_file, exist_ok=True)
|
|
name = self.model_name.replace('/', '-')
|
|
output_file = os.path.join(output_file, name)
|
|
if model_type in ['pytorch', 'torchscript']:
|
|
output_file = output_file + '.pt'
|
|
elif model_type == 'onnx':
|
|
output_file = output_file + '.onnx'
|
|
else:
|
|
raise AttributeError('Unsupported model_type.')
|
|
|
|
dummy_input = 'test sentence'
|
|
inputs = self.tokenizer(dummy_input, padding=True, truncation=True, return_tensors='pt') # a dictionary
|
|
if model_type == 'pytorch':
|
|
torch.save(self._model, output_file)
|
|
elif model_type == 'torchscript':
|
|
inputs = list(inputs.values())
|
|
try:
|
|
try:
|
|
jit_model = torch.jit.script(self._model)
|
|
except Exception:
|
|
jit_model = torch.jit.trace(self._model, inputs, strict=False)
|
|
torch.jit.save(jit_model, output_file)
|
|
except Exception as e:
|
|
log.error(f'Fail to save as torchscript: {e}.')
|
|
raise RuntimeError(f'Fail to save as torchscript: {e}.')
|
|
elif model_type == 'onnx':
|
|
dynamic_axes = {}
|
|
for k, v in self.onnx_config['inputs'].items():
|
|
dynamic_axes[k] = v
|
|
for k, v in self.onnx_config['outputs'].items():
|
|
dynamic_axes[k] = v
|
|
torch.onnx.export(
|
|
self._model,
|
|
tuple(inputs.values()),
|
|
output_file,
|
|
input_names=list(self.onnx_config['inputs'].keys()),
|
|
output_names=list(self.onnx_config['outputs'].keys()),
|
|
dynamic_axes=dynamic_axes,
|
|
opset_version=torch.onnx.constant_folding_opset_versions[-1] if hasattr(
|
|
torch.onnx, 'constant_folding_opset_versions') else 14,
|
|
do_constant_folding=True,
|
|
)
|
|
# todo: elif format == 'tensorrt':
|
|
else:
|
|
log.error(f'Unsupported format "{format}".')
|
|
return Path(output_file).resolve()
|
|
|
|
@property
|
|
def supported_formats(self):
|
|
onnxes = self.supported_model_names(format='onnx')
|
|
if self.model_name in onnxes:
|
|
return ['onnx']
|
|
else:
|
|
return ['pytorch']
|
|
|
|
@property
|
|
def supported_formats(self):
|
|
return ['onnx']
|
|
|
|
@staticmethod
|
|
def supported_model_names(format: str = None):
|
|
full_list = [
|
|
|
|
]
|
|
full_list.sort()
|
|
if format is None:
|
|
model_list = full_list
|
|
elif format == 'pytorch':
|
|
to_remove = []
|
|
assert set(to_remove).issubset(set(full_list))
|
|
model_list = list(set(full_list) - set(to_remove))
|
|
elif format == 'torchscript':
|
|
to_remove = [
|
|
]
|
|
assert set(to_remove).issubset(set(full_list))
|
|
model_list = list(set(full_list) - set(to_remove))
|
|
elif format == 'onnx':
|
|
to_remove = [
|
|
]
|
|
assert set(to_remove).issubset(set(full_list))
|
|
model_list = list(set(full_list) - set(to_remove))
|
|
# todo: elif format == 'tensorrt':
|
|
else:
|
|
log.error(f'Invalid format "{format}". Currently supported formats: "pytorch", "torchscript".')
|
|
return model_list
|
|
|
|
@staticmethod
|
|
def map_model_names(name):
|
|
req = requests.get("https://www.sbert.net/_static/html/models_en_sentence_embeddings.html")
|
|
data = req.text
|
|
default_sbert = []
|
|
for line in data.split('\r\n'):
|
|
line = line.replace(' ', '')
|
|
if line.startswith('"name":'):
|
|
n = line.split(':')[-1].replace('"', '').replace(',', '')
|
|
default_sbert.append(n)
|
|
if name in default_sbert:
|
|
name = 'sentence-transformers/' + name
|
|
return name
|