def __init__(self,
model_path: Optional[str] = 'pretrained',
channel_axis: int = 1,
on_gpu: bool = False,
*args,
**kwargs):
super().__init__(*args, **kwargs)
self.channel_axis = channel_axis
self.model_path = model_path
self.on_gpu = on_gpu
self.logger = logger
if self.model_path and os.path.exists(self.model_path):
import tensorflow as tf
cpus = tf.config.experimental.list_physical_devices(
device_type='CPU')
gpus = tf.config.experimental.list_physical_devices(
device_type='GPU')
if self.on_gpu and len(gpus) > 0:
cpus.append(gpus[0])
tf.config.experimental.set_visible_devices(devices=cpus)
self.logger.info(f'BiT model path: {self.model_path}')
from tensorflow.python.keras.models import load_model
_model = load_model(self.model_path)
self.model = _model.signatures['serving_default']
self._get_input = tf.convert_to_tensor
else:
raise PretrainedModelFileDoesNotExist(
f'model at {self.model_path} does not exist')
def post_init(self):
"""
Load the model from the `.onnx` file and add outputs for the selected layer, i.e. ``outputs_name``. The modified
models is saved at `tmp_model_path`.
"""
super().post_init()
model_name = self.raw_model_path.split(
'/')[-1] if self.raw_model_path else None
tmp_model_path = self.get_file_from_workspace(
f'{model_name}.tmp') if model_name else None
raw_model_path = self.raw_model_path
if self.raw_model_path and is_url(self.raw_model_path):
import urllib.request
download_path, *_ = urllib.request.urlretrieve(self.raw_model_path)
raw_model_path = download_path
self.logger.info(f'download the model at {self.raw_model_path}')
if tmp_model_path and not os.path.exists(
tmp_model_path) and self.outputs_name:
self._append_outputs(raw_model_path, self.outputs_name,
tmp_model_path)
self.logger.info(
f'save the model with outputs [{self.outputs_name}] at {tmp_model_path}'
)
if tmp_model_path and os.path.exists(tmp_model_path):
import onnxruntime
self.model = onnxruntime.InferenceSession(tmp_model_path, None)
self.inputs_name = self.model.get_inputs()[0].name
self._device = None
self.to_device(self.model)
else:
raise PretrainedModelFileDoesNotExist(
f'model at {tmp_model_path} does not exist')
def post_init(self):
self.to_device()
if self.model_path and os.path.exists(self.model_path):
import tensorflow as tf
model = tf.keras.models.load_model(self.model_path)
model.trainable = False
self.model = tf.keras.Model(inputs=model.input,
outputs=model.get_layer(self.layer_name).output)
else:
raise PretrainedModelFileDoesNotExist(f'model {self.model_path} does not exist')
def post_init(self):
super().post_init()
if self.model_path and os.path.exists(self.model_path):
self.to_device()
import tensorflow as tf
self.logger.info(f'model_path: {self.model_path}')
_model = tf.saved_model.load(self.model_path)
self.model = _model.signatures['serving_default']
self._get_input = tf.convert_to_tensor
else:
raise PretrainedModelFileDoesNotExist(f'model at {self.model_path} does not exist')
def post_init(self):
import os
import pickle
super().post_init()
if os.path.exists(self.path_vectorizer):
self.tfidf_vectorizer = pickle.load(
open(self.path_vectorizer, 'rb'))
else:
raise PretrainedModelFileDoesNotExist(
f'{self.path_vectorizer} not found, cannot find a fitted tfidf_vectorizer'
)
def post_init(self):
"""
Load the model from the `.ckpt` checkpoint.
"""
from mindspore.train.serialization import load_checkpoint, load_param_into_net
super().post_init()
if self.model_path and os.path.exists(self.model_path):
self.to_device()
_param_dict = load_checkpoint(ckpt_file_name=self.model_path)
load_param_into_net(self.model, _param_dict)
else:
raise PretrainedModelFileDoesNotExist(f'model {self.model_path} does not exist')
def post_init(self):
super().post_init()
import pickle
from .pyngramspell import PyNgramSpell, BKTree
self.model = None
if os.path.exists(self.model_path):
with open(self.model_path, 'rb') as model_file:
self.model = pickle.load(model_file)
else:
raise PretrainedModelFileDoesNotExist(
f'{self.model_path} not found, cannot find a fitted spell checker'
)
def post_init(self):
super().post_init()
import torch
if self.model_path and os.path.exists(self.model_path):
if self.pool_strategy is not None:
self.pool_fn = getattr(np, self.pool_strategy)
self.model = torch.load(self.model_path)
self.model.eval()
self.to_device(self.model)
self.layer = getattr(self.model, self.layer_name)
else:
raise PretrainedModelFileDoesNotExist(
f'model {self.model_path} does not exist')
def post_init(self):
super().post_init()
import torch
if self.model_path and os.path.exists(self.model_path):
with open(self.texts_path, 'rb') as fp:
texts = pickle.load(fp)
self.model = TIRG(texts, 512)
model_sd = torch.load(self.model_path, map_location=torch.device('cpu'))
self.model.load_state_dict(model_sd['model_state_dict'])
self.model.eval()
self.to_device(self.model)
else:
raise PretrainedModelFileDoesNotExist(f'model {self.model_path} does not exist')
def post_init(self):
super().post_init()
if self.model_path and os.path.exists(self.model_path):
import lightgbm
self.booster = lightgbm.Booster(model_file=self.model_path)
model_num_features = self.booster.num_feature()
expected_num_features = len(self.query_feature_names + self.match_feature_names)
if model_num_features != expected_num_features:
raise ValueError(
f'The number of features expected by the LightGBM model {model_num_features} is different'
f'than the ones provided in input {expected_num_features}')
else:
raise PretrainedModelFileDoesNotExist(f'model {self.model_path} does not exist')
def post_init(self):
super().post_init()
if self.model_path and os.path.exists(self.model_path):
import torch
from fairseq.models.wav2vec import Wav2VecModel
cp = torch.load(self.model_path, map_location=torch.device('cpu'))
self.model = Wav2VecModel.build_model(cp['args'], task=None)
self.model.load_state_dict(cp['model'])
self.model.eval()
self.to_device(self.model)
self._tensor_func = torch.tensor
else:
raise PretrainedModelFileDoesNotExist(
f'model at {self.model_path} does not exist')
def post_init(self):
"""Load VAE model"""
super().post_init()
from cvae import cvae
import cvae.lib.model_iaf as model
import tensorflow as tf
params_path = os.path.join(
self.model_path,
'params.json') if self.model_path and os.path.exists(
self.model_path) else None
if params_path and os.path.exists(params_path):
config = tf.ConfigProto(log_device_placement=False)
config.gpu_options.allow_growth = True
with tf.Graph().as_default():
# Load parameter file.
with open(params_path, 'r') as f:
param = json.load(f)
net = model.VAEModel(param,
None,
input_dim=param['dim_feature'],
keep_prob=tf.placeholder_with_default(
input=tf.cast(1.0, dtype=tf.float32),
shape=(),
name="KeepProb"),
initializer='orthogonal')
# Placeholder for data features
self.data_feature_placeholder = tf.placeholder_with_default(
input=tf.zeros([64, param['dim_feature']],
dtype=tf.float32),
shape=[None, param['dim_feature']])
self.embeddings = net.embed(self.data_feature_placeholder)
self.sess = tf.Session(config=config)
init = tf.global_variables_initializer()
self.sess.run(init)
# Saver for loading checkpoints of the model.
saver = tf.train.Saver(var_list=tf.trainable_variables())
cvae.load(saver, self.sess, self.model_path)
self.to_device()
else:
raise PretrainedModelFileDoesNotExist()
