def smoke_test_model(model_path):
try:
resolved_model = hub.resolve(model_path)
loader_impl.parse_saved_model(resolved_model)
except Exception as e: # pylint: disable=broad-except
return False, e
return True, None
def build_encoding(encoding_strat, text_input):
if encoding_strat == 'tfidf':
encoder = tf.keras.layers.experimental.preprocessing.TextVectorization(
max_tokens=VOCAB_SIZE, output_mode='tf-idf')
encoder.adapt(train_dataset.map(lambda text, label: text))
net = encoder(text_input)
return net
elif encoding_strat == 'learned':
encoder = tf.keras.layers.experimental.preprocessing.TextVectorization(
max_tokens=VOCAB_SIZE, output_mode='int')
encoder.adapt(train_dataset.map(lambda text, label: text))
net = encoder(text_input)
net = tf.keras.layers.Embedding(input_dim=len(
encoder.get_vocabulary()),
output_dim=64,
mask_zero=True)(net)
net = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(10,
name='LSTM'),
name='Bidirectional')(net)
return net
elif encoding_strat == 'bert':
load_options = tf.saved_model.LoadOptions(
experimental_io_device='/job:localhost')
preprocessor = hub.load(
"https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3",
options=load_options)
print("\n\n\nHub resolve:")
print(
hub.resolve(
"https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3"))
print("\n\n")
# Step 1: tokenize batches of text inputs.
tokenize = hub.KerasLayer(preprocessor.tokenize, name='tokenizer')
tokenized_inputs = [
tokenize(text_input),
]
# Step 2 (optional): modify tokenized inputs.
# Step 3: pack input sequences for the Transformer encoder.
seq_length = 250 # We filter out anything more earlier
bert_pack_inputs = hub.KerasLayer(
preprocessor.bert_pack_inputs,
arguments=dict(seq_length=seq_length),
name='input_packer') # Optional argument.
encoder_inputs = bert_pack_inputs(tokenized_inputs)
encoder = hub.KerasLayer(
"https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-6_H-512_A-8/1",
trainable=False,
name='BERT_encoder')
outputs = encoder(encoder_inputs)
sequence_output = outputs["sequence_output"]
net = tf.keras.layers.Bidirectional(
tf.keras.layers.LSTM(10, name='LSTM'),
name='Bidirectional')(sequence_output)
return net
else:
raise ValueError(f"Invalid encoding strategy - given {encoding_strat}")
def __init__(self, num_units, CLS, SEP, PAD, UNK, min_len, max_len,
bert_hub_url, **kwargs):
""" Initializes the layer
For more details on parameters, check the argument parser in run_detext.py
"""
super(BertLayer, self).__init__()
self.text_ftr_size = num_units
# Load pretrained model using hub
with tf.name_scope(BERT_VAR_PREFIX):
self.text_encoder = hub.KerasLayer(hub.resolve(bert_hub_url),
trainable=True)
self.last_layer_cls = tf.keras.layers.Dense(num_units,
activation='tanh',
use_bias=True,
name=BERT_VAR_PREFIX +
'_last_layer_cls')
self.preprocess_layer = BertPreprocessLayer(self.text_encoder, max_len,
min_len, CLS, SEP, PAD,
UNK)
self._pad_id = self.preprocess_layer.pad_id()
self._pad_id = tf.cast(self._pad_id, tf.int32)
self.dropout = tf.keras.layers.Dropout(0.1)
def convert_tf_hub_module(module_handle, output_dir,
signature='default', saved_model_tags='serve',
quantization_dtype_map=None,
skip_op_check=False, strip_debug_ops=False,
weight_shard_size_bytes=1024 * 1024 * 4,
control_flow_v2=False,
experiments=False,
metadata=None):
"""Conversion for TF Hub modules V1 and V2.
See convert_tf_hub_module and convert_tf_saved_model.
Args:
module_path: string Path to the module.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'model.json'
- possibly sharded binary weight files.
signature: string Signature to load.
saved_model_tags: tags of the GraphDef to load. Defaults to ''.
quantization_dtype_map: A mapping from dtype
(`uint8`, `uint16`, `float16`) to weights names. The weight mapping
supports wildcard substitution.
skip_op_check: Bool whether to skip the op check.
strip_debug_ops: Bool whether to strip debug ops.
weight_shard_size_bytes: Shard size (in bytes) of the weight files.
The size of each weight file will be <= this value.
control_flow_v2: Bool whether to enable control flow v2 ops.
experiments: Bool enable experimental features.
metadata: User defined metadata map.
"""
module_path = hub.resolve(module_handle)
# TODO(vbardiovskyg): We can remove this v1 code path once loading of all v1
# modules is fixed on the TF side, or once the modules we cannot load become
# replaced with newer versions.
if tf.io.gfile.exists(os.path.join(module_path, _HUB_V1_MODULE_PB)):
print("Loading the module using TF 1.X interface from %s." % module_path)
convert_tf_hub_module_v1(module_path, output_dir, signature,
quantization_dtype_map,
skip_op_check, strip_debug_ops,
weight_shard_size_bytes,
experiments=experiments,
metadata=metadata)
else:
print("Loading the module using TF 2.X interface from %s." % module_path)
if signature is None:
signature = 'default'
convert_tf_saved_model(saved_model_dir=module_path,
output_dir=output_dir,
signature_def=signature,
saved_model_tags=saved_model_tags,
quantization_dtype_map=quantization_dtype_map,
skip_op_check=skip_op_check,
strip_debug_ops=strip_debug_ops,
weight_shard_size_bytes=weight_shard_size_bytes,
control_flow_v2=control_flow_v2,
experiments=experiments,
metadata=metadata)
def test_resolve(self):
with tf.Graph().as_default():
self._generate_module()
module_dir = hub.resolve(
"http://localhost:%d/test_module.tgz" % self.server_port)
self.assertIn(tempfile.gettempdir(), module_dir)
module_files = sorted(tf_v1.gfile.ListDirectory(module_dir))
self.assertEqual(
["assets", "saved_model.pb", "tfhub_module.pb", "variables"],
module_files)
def assert_can_resolve_asset(self, asset_path: str):
"""Attempt to hub.resolve the given asset path."""
try:
resolved_model = hub.resolve(asset_path)
loader_impl.parse_saved_model(resolved_model)
_validate_file_paths(resolved_model)
except Exception as e: # pylint: disable=broad-except
raise MarkdownDocumentationError(
f"The model on path {asset_path} failed to parse. Please make sure "
"that the asset-path metadata points to a valid TF2 SavedModel or a "
"TF1 Hub module, compressed as described in section 'Model' of "
f"README.md. Underlying reason for failure: {e}.")
def __init__(self, tfhub_handle = DEFAULT_TOKENIZER_TFHUB_HANDLE):
preprocessor = hub.load(tfhub_handle)
model_path = hub.resolve(tfhub_handle)
vocab_file_path = os.path.join(model_path, 'assets/vocab.txt')
with tf.io.gfile.GFile(vocab_file_path, 'r') as vocab_file:
self.vocab = [token.strip() for token in vocab_file.readlines()]
text_input = tf.keras.layers.Input(shape=(), dtype=tf.string)
tokenize_layer = hub.KerasLayer(preprocessor.tokenize)
outputs = tokenize_layer(text_input)
self.model = tf.keras.Model(
inputs=text_input, outputs=outputs, name='tokenizer')
def assert_can_resolve_asset(self, asset_path: str):
"""Attempt to hub.resolve the given asset path."""
try:
resolved_model = hub.resolve(asset_path)
loader_impl.parse_saved_model(resolved_model)
_validate_file_paths(resolved_model)
except Exception as e: # pylint: disable=broad-except
raise MarkdownDocumentationError(
f"The model on path {asset_path} failed to parse. Please make sure "
"that the asset-path metadata points to a valid TF2 SavedModel or a "
"TF1 Hub module as described on "
"https://www.tensorflow.org/hub/exporting_tf2_saved_model. "
f"Underlying reason for failure: {e}.")
def convert_tf_hub_module(module_handle,
output_dir,
signature='default',
saved_model_tags='serve',
quantization_dtype=None,
skip_op_check=False,
strip_debug_ops=False,
weight_shard_size_bytes=1024 * 1024 * 4):
"""Conversion for TF Hub modules V1 and V2.
See convert_tf_hub_module and convert_tf_saved_model.
Args:
module_path: string Path to the module.
output_dir: string The name of the output directory. The directory
will consist of
- a file named 'model.json'
- possibly sharded binary weight files.
signature: string Signature to load.
saved_model_tags: tags of the GraphDef to load. Defaults to ''.
skip_op_check: Bool whether to skip the op check.
strip_debug_ops: Bool whether to strip debug ops.
weight_shard_size_bytes: Shard size (in bytes) of the weight files.
The size of each weight file will be <= this value.
"""
module_path = hub.resolve(module_handle)
# TODO(vbardiovskyg): We can remove this v1 code path once loading of all v1
# modules is fixed on the TF side, or once the modules we cannot load become
# replaced with newer versions.
if tf.io.gfile.exists(os.path.join(module_path, _HUB_V1_MODULE_PB)):
print("Loading the module using TF 1.X interface from %s." %
module_path)
convert_tf_hub_module_v1(module_path, output_dir, signature,
quantization_dtype, skip_op_check,
strip_debug_ops, weight_shard_size_bytes)
else:
print("Loading the module using TF 2.X interface from %s." %
module_path)
if signature is None:
signature = 'default'
convert_tf_saved_model(saved_model_dir=module_path,
output_dir=output_dir,
signature_def=signature,
saved_model_tags=saved_model_tags,
quantization_dtype=quantization_dtype,
skip_op_check=skip_op_check,
strip_debug_ops=strip_debug_ops,
weight_shard_size_bytes=weight_shard_size_bytes)
def import_from_tfhub(
identifier: t.Union[str, "HubModule"],
name: t.Optional[str] = None,
labels: t.Optional[t.Dict[str, str]] = None,
custom_objects: t.Optional[t.Dict[str, t.Any]] = None,
metadata: t.Optional[t.Dict[str, t.Any]] = None,
model_store: "ModelStore" = Provide[BentoMLContainer.model_store],
) -> Tag:
"""
Import a model from `Tensorflow Hub <https://tfhub.dev/>`_ to BentoML modelstore.
Args:
identifier (:code:`Union[str, tensorflow_hub.Module, tensorflow_hub.KerasLayer]`): Identifier accepts
two type of inputs:
- if `type` of :code:`identifier` either of type :code:`tensorflow_hub.Module` (**legacy** `tensorflow_hub`) or :code:`tensorflow_hub.KerasLayer` (`tensorflow_hub`), then we will save the given model to a :code:`SavedModel` format.
- if `type` of :code:`identifier` is a :obj:`str`, we assume that this is the URI retrieved from Tensorflow Hub. We then clean the given URI, and get a local copy of a given model to BentoML modelstore. name (:code:`str`, `optional`, defaults to `None`): An optional name for the model. If :code:`identifier` is a :obj:`str`, then name can be autogenerated from the given URI.
name (:code:`str`, `optional`, default to `None`):
Optional name for the saved model. If None, then name will be generated from :code:`identifier`.
labels (:code:`Dict[str, str]`, `optional`, default to :code:`None`):
user-defined labels for managing models, e.g. team=nlp, stage=dev
custom_objects (:code:`Dict[str, Any]]`, `optional`, default to :code:`None`):
user-defined additional python objects to be saved alongside the model,
e.g. a tokenizer instance, preprocessor function, model configuration json
metadata (:code:`Dict[str, Any]`, `optional`, default to :code:`None`):
Custom metadata for given model.
model_store (:mod:`~bentoml._internal.models.store.ModelStore`, default to :mod:`BentoMLContainer.model_store`):
BentoML modelstore, provided by DI Container.
Returns:
:obj:`~bentoml.Tag`: A :obj:`~bentoml.Tag` object that can be used to retrieve the model with :func:`bentoml.tensorflow.load`:
Example for importing a model from Tensorflow Hub:
.. code-block:: python
import tensorflow_text as text # noqa # pylint: disable
import bentoml
tag = bentoml.tensorflow_v1.import_from_tfhub("https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3")
# load model back with `load`:
model = bentoml.tensorflow_v1.load(tag, load_as_hub_module=True)
Example for importing a custom Tensorflow Hub model:
.. code-block:: python
import tensorflow as tf
import tensorflow_hub as hub
import bentoml
# Simple toy Tensorflow Hub model
def _plus_one_model_tf1():
def plus_one():
x = tf.placeholder(dtype=tf.float32, name="x")
y = x + 1
hub.add_signature(inputs=x, outputs=y)
spec = hub.create_module_spec(plus_one)
with tf.get_default_graph().as_default():
module = hub.Module(spec, trainable=True)
return module
model = _plus_one_model_tf1()
# retrieve the given tag:
tag = bentoml.tensorflow_v1.import_from_tfhub(model)
""" # noqa
if hub is None:
raise MissingDependencyException("""\
`tensorflow_hub` does not exists.
Make sure to `pip install --upgrade tensorflow_hub` before using.
""")
context: t.Dict[str, t.Any] = {
"framework_name":
"tensorflow",
"pip_dependencies": [
f"tensorflow=={get_tf_version()}",
f"tensorflow_hub=={importlib_metadata.version('tensorflow_hub')}",
],
"import_from_tfhub":
True,
}
if name is None:
if isinstance(identifier, str):
name = _clean_name(identifier)
else:
name = f"{identifier.__class__.__name__}_{uuid.uuid4().hex[:5].upper()}"
with bentoml.models.create(
name,
module=MODULE_NAME,
options=None,
context=context,
labels=labels,
custom_objects=custom_objects,
metadata=metadata,
) as _model:
if isinstance(identifier, str):
current_cache_dir = os.environ.get("TFHUB_CACHE_DIR")
os.environ["TFHUB_CACHE_DIR"] = _model.path
fpath: str = resolve(identifier)
folder = fpath.split("/")[-1]
_model.info.options = {"model": identifier, "local_path": folder}
if current_cache_dir is not None:
os.environ["TFHUB_CACHE_DIR"] = current_cache_dir
else:
if hasattr(identifier, "export"):
# hub.Module.export()
with tf.Session(
graph=tf.get_default_graph()) as sess: # type: ignore
sess.run(tf.global_variables_initializer()) # type: ignore
identifier.export(_model.path, sess) # type: ignore
else:
tf.saved_model.save(identifier, _model.path)
_model.info.options = {
"model": identifier.__class__.__name__,
"local_path": ".",
}
return _model.tag
class TestBertLayer(tf.test.TestCase, DataSetup):
"""Unit test for bert_layer.py"""
# Bert setup
bert_hub_layer = hub.KerasLayer(hub.resolve(DataSetup.bert_hub_url),
trainable=True)
bert_vocab_file = bert_hub_layer.resolved_object.vocab_file.asset_path.numpy(
).decode("utf-8")
bert_vocab_table = read_vocab(bert_vocab_file)
bert_PAD_ID = bert_vocab_table[DataSetup.PAD]
bert_SEP_ID = bert_vocab_table[DataSetup.SEP]
bert_CLS_ID = bert_vocab_table[DataSetup.CLS]
bert_UNK_ID = bert_vocab_table[DataSetup.UNK]
# SentencePiece setup
sentencepiece_hub_layer = hub.KerasLayer(
hub.resolve(DataSetup.libert_sp_hub_url))
tokenizer_file = sentencepiece_hub_layer.resolved_object.tokenizer_file.asset_path.numpy(
).decode("utf-8")
with tf.io.gfile.GFile(tokenizer_file, 'rb') as f_handler:
sp_model = f_handler.read()
sentencepiece_tokenizer = tf_text.SentencepieceTokenizer(model=sp_model,
out_type=tf.int32)
sentencepiece_vocab_tf_table = create_tf_vocab_from_sp_tokenizer(
sp_tokenizer=sentencepiece_tokenizer, num_oov_buckets=1)
sentencepiece_PAD_ID = sentencepiece_vocab_tf_table.lookup(
tf.constant(DataSetup.PAD)).numpy()
sentencepiece_SEP_ID = sentencepiece_vocab_tf_table.lookup(
tf.constant(DataSetup.SEP)).numpy()
sentencepiece_CLS_ID = sentencepiece_vocab_tf_table.lookup(
tf.constant(DataSetup.CLS)).numpy()
sentencepiece_UNK_ID = sentencepiece_vocab_tf_table.lookup(
tf.constant(DataSetup.UNK)).numpy()
# Space setup
space_hub_layer = hub.KerasLayer(
hub.resolve(DataSetup.libert_space_hub_url))
tokenizer_file = space_hub_layer.resolved_object.tokenizer_file.asset_path.numpy(
).decode("utf-8")
space_vocab = read_vocab(tokenizer_file)
space_PAD_ID = space_vocab[DataSetup.PAD]
space_SEP_ID = space_vocab[DataSetup.SEP]
space_CLS_ID = space_vocab[DataSetup.CLS]
space_UNK_ID = space_vocab[DataSetup.UNK]
# Hyperparameters setup
num_units = 16
pad_id = 0
num_doc_fields = 2
min_len = 3
max_len = 8
layer = bert_layer.BertLayer(num_units, DataSetup.CLS, DataSetup.SEP,
DataSetup.PAD, DataSetup.UNK, min_len,
max_len, DataSetup.bert_hub_url)
def testBertLayer(self):
"""Test Bert layer """
for hub_url in [self.bert_hub_url]:
self._testBertLayer(hub_url)
def _testBertLayer(self, hub_url):
query = self.query
doc_fields = [self.ranking_doc_field1, self.ranking_doc_field2]
user_fields = [query, query, query]
query_ftrs, doc_ftrs, user_ftrs = self.layer(
{
InputFtrType.QUERY_COLUMN_NAME: query,
InputFtrType.DOC_TEXT_COLUMN_NAMES: doc_fields,
InputFtrType.USER_TEXT_COLUMN_NAMES: user_fields
}, False)
text_ftr_size = self.num_units
self.assertEqual(text_ftr_size, self.layer.text_ftr_size)
self.assertAllEqual(query_ftrs.shape, [2, self.layer.text_ftr_size])
self.assertAllEqual(doc_ftrs.shape,
[2, 3, 2, self.layer.text_ftr_size])
self.assertAllEqual(user_ftrs.shape, [2, 3, self.layer.text_ftr_size])
# 1st query, 2nd doc, 2nd field should be the same as 2nd query, 1st doc, 2nd field
self.assertAllEqual(doc_ftrs[0, 1, 1], doc_ftrs[1, 0, 1])
# 1st query, 1st doc, 1st field should be the same as 1st query, 1st doc, 2nd field
self.assertAllEqual(doc_ftrs[0, 0, 0], doc_ftrs[0, 0, 1])
# 1st query, 1st doc, 2st field should NOT be the same as 1st query, 2st doc, 2nd field
self.assertNotAllClose(doc_ftrs[0, 1, 0], doc_ftrs[0, 1, 1])
def testGetInputIds(self):
"""Tests get_input_ids() """
query = [[1, 2, 3], [2, 4, 3]]
doc_field1 = [[[1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0]],
[[2, 4, 3, 1], [1, 3, 3, 1], [1, 3, 3, 1]]]
doc_field2 = [[[1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0]],
[[20, 5, 3, 1], [5, 6, 1, 1], [5, 6, 1, 1]]]
query = tf.constant(query, dtype=tf.int32)
doc_field1 = tf.constant(doc_field1, dtype=tf.int32)
doc_field2 = tf.constant(doc_field2, dtype=tf.int32)
doc_fields = [doc_field1, doc_field2]
user_fields = None
max_text_len, max_text_len_array = bert_layer.BertLayer.get_input_max_len(
query, doc_fields, user_fields)
bert_input_ids = bert_layer.BertLayer.get_bert_input_ids(
query, doc_fields, user_fields, self.pad_id, max_text_len,
max_text_len_array)
# Check bert input ids
self.assertAllEqual(
bert_input_ids,
[[1, 2, 3, 0], [2, 4, 3, 0], [1, 2, 3, 0], [2, 4, 3, 1],
[0, 0, 0, 0], [2, 4, 3, 1], [1, 3, 3, 1], [1, 3, 3, 1],
[1, 2, 3, 0], [2, 4, 3, 1], [0, 0, 0, 0], [20, 5, 3, 1],
[5, 6, 1, 1], [5, 6, 1, 1]])
def testPrerocessQuery(self):
"""Tests _preprocess_query function of bert layer"""
query = tf.constant(['batch 1 user 1 build', 'batch 2 user 2 word'],
dtype=tf.string)
expected = tf.constant(
[[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 4, 2
],
[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 5, 2
]],
dtype=tf.int32)
self.assertAllEqual(expected, self.layer._preprocess_query(query))
def testPrerocessUsr(self):
"""Tests _preprocess_user function of bert layer"""
user_fields = [
tf.constant(['batch 1 user 1 build', 'batch 2 user 2 word'],
dtype=tf.string)
]
expected = [
tf.constant([[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 4, 2
],
[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 5, 2
]],
dtype=tf.int32)
]
self.assertAllEqual(expected, self.layer._preprocess_user(user_fields))
def testPrerocessDoc(self):
"""Tests _preprocess_doc function of bert layer"""
doc_fields = [
tf.constant([['batch 1 doc 1 build', 'batch 1 doc 2'],
['batch 2 doc 1 word', 'batch 2 doc 2']],
dtype=tf.string)
]
expected = [
tf.constant([[[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 4, 2
],
[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 2, 3
]],
[[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 5, 2
],
[
self.CLS_ID, self.bert_UNK_ID, self.bert_UNK_ID,
self.bert_UNK_ID, self.bert_UNK_ID, 2, 3
]]],
dtype=tf.int32)
]
self.assertAllEqual(expected, self.layer._preprocess_doc(doc_fields))
def testBertPreprocessLayerWordPiece(self):
"""Tests BertPreprocessLayer with wordpiece tokenizer"""
preprocess_layer = bert_layer.BertPreprocessLayer(
self.bert_hub_layer, self.max_len, self.min_len, self.CLS,
self.SEP, self.PAD, self.UNK)
sentences = tf.constant(
['test sent1', 'build build build build sent2'])
expected = tf.constant([[
self.bert_CLS_ID, 8, self.bert_UNK_ID, self.bert_SEP_ID,
self.bert_PAD_ID, self.bert_PAD_ID, self.bert_PAD_ID
], [self.bert_CLS_ID, 4, 4, 4, 4, self.bert_UNK_ID, self.bert_SEP_ID]],
dtype=tf.int32)
outputs = preprocess_layer(sentences)
self.assertAllEqual(expected, outputs)
def testBertPreprocessLayerSentencePiece(self):
"""Tests BertPreprocessLayer with sentencepiece tokenizer"""
preprocess_layer = bert_layer.BertPreprocessLayer(
self.sentencepiece_hub_layer, self.max_len, self.min_len, self.CLS,
self.SEP, self.PAD, self.UNK)
sentences = tf.constant(
['TEST sent1', 'build build build build sent2'])
expected = tf.constant([[
self.sentencepiece_CLS_ID, 557, 4120, 29900,
self.sentencepiece_SEP_ID, self.sentencepiece_PAD_ID,
self.sentencepiece_PAD_ID, self.sentencepiece_PAD_ID
],
[
self.sentencepiece_CLS_ID, 671, 671, 671,
671, 4120, 29904, self.sentencepiece_SEP_ID
]],
dtype=tf.int32)
outputs = preprocess_layer(sentences)
self.assertAllEqual(expected, outputs)
def testBertPreprocessLayerSpace(self):
"""Tests BertPreprocessLayer with space tokenizer"""
preprocess_layer = bert_layer.BertPreprocessLayer(
self.space_hub_layer, self.max_len, self.min_len, self.CLS,
self.SEP, self.PAD, self.UNK)
sentences = tf.constant(
['test sent1', 'build build build build sent2'])
expected = tf.constant([[
self.space_CLS_ID, 8, self.space_UNK_ID, self.space_SEP_ID,
self.space_PAD_ID, self.space_PAD_ID, self.space_PAD_ID
], [
self.space_CLS_ID, 4, 4, 4, 4, self.space_UNK_ID, self.space_SEP_ID
]],
dtype=tf.int32)
outputs = preprocess_layer(sentences)
self.assertAllEqual(expected, outputs)
def testBertPreprocessLayerAdjustLen(self):
"""Tests adjust_len function of BertPreprocessLayer"""
sentences = tf.constant(
['test sent1', 'build build build build sent2'])
min_len = 12
max_len = 16
layer = bert_layer.BertPreprocessLayer(self.bert_hub_layer, max_len,
min_len, self.CLS, self.SEP,
self.PAD, self.UNK)
outputs = layer(sentences)
shape = tf.shape(outputs)
self.assertAllEqual(shape, tf.constant([2, 12]))
min_len = 0
max_len = 1
layer = bert_layer.BertPreprocessLayer(self.bert_hub_layer, max_len,
min_len, self.CLS, self.SEP,
self.PAD, self.UNK)
outputs = layer(sentences)
shape = tf.shape(outputs)
self.assertAllEqual(shape, tf.constant([2, 1]))
def model_fn(features, mode, params):
"""A function for applying hub module that follows Estimator API."""
hub_module = params["hub_module"]
finetune_layer = params["finetune_layer"]
num_classes = params["num_classes"]
initial_learning_rate = params["initial_learning_rate"]
momentum = params["momentum"]
lr_decay_factor = params["lr_decay_factor"]
decay_steps = params["decay_steps"]
warmup_steps = params["warmup_steps"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
module_path = hub.resolve(hub_module)
is_legacy_hub_module = tf.io.gfile.exists(
os.path.join(module_path, "tfhub_module.pb"))
if is_legacy_hub_module:
module = hub.Module(hub_module,
trainable=is_training,
tags={"train"} if is_training else None)
pre_logits = module(features["image"],
signature=params["hub_module_signature"],
as_dict=True)[finetune_layer]
else:
module = hub.load(hub_module)
tf.get_collection_ref(tf.GraphKeys.TRAINABLE_VARIABLES).extend(
module.trainable_variables)
pre_logits = module(features["image"], training=is_training)
num_dim_pre_logits = len(pre_logits.get_shape().as_list())
if num_dim_pre_logits == 4:
pre_logits = tf.squeeze(pre_logits, [1, 2])
elif num_dim_pre_logits != 2:
raise ValueError("Invalid number of dimensions in the representation "
"layer: {}, but only 2 or 4 are allowed".format(
num_dim_pre_logits))
logits = tf.layers.dense(pre_logits,
units=num_classes,
kernel_initializer=tf.zeros_initializer(),
name="linear_head")
loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=logits, labels=features["label"])
loss = tf.reduce_mean(loss)
def accuracy_metric(logits, labels):
return {"accuracy": tf.metrics.accuracy(
labels=labels,
predictions=tf.argmax(logits, axis=-1))}
eval_metrics = (accuracy_metric, [logits, features["label"]])
if mode == tf.estimator.ModeKeys.EVAL:
if params["tpu_name"] is not None:
return tf.contrib.tpu.TPUEstimatorSpec(
mode=mode, loss=loss, eval_metrics=eval_metrics)
else:
return tf.estimator.EstimatorSpec(
mode=mode, loss=loss,
eval_metric_ops=eval_metrics[0](*eval_metrics[1]))
elif mode == tf.estimator.ModeKeys.TRAIN:
train_op = trainer.get_train_op(loss,
initial_learning_rate,
momentum,
lr_decay_factor,
decay_steps,
warmup_steps,
use_tpu=params["tpu_name"] is not None)
spec_type = (tf.contrib.tpu.TPUEstimatorSpec
if params["tpu_name"] is not None
else tf.estimator.EstimatorSpec)
return spec_type(mode=mode, loss=loss, train_op=train_op)
else:
raise ValueError("Unsupported mode %s" % mode)
