def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
temperature=1.0,
sentencepiece_model_path=t5.data.DEFAULT_SPM_PATH):
"""Exports a TensorFlow SavedModel.
Args:
export_dir: str, a directory in which to export SavedModels. Will use
`model_dir` if unspecified.
checkpoint_step: int, checkpoint to export. If -1 (default), use the
latest checkpoint from the pretrained model directory.
beam_size: int, a number >= 1 specifying the number of beams to use for
beam search.
temperature: float, a value between 0 and 1 (must be 0 if beam_size > 1)
0.0 means argmax, 1.0 means sample according to predicted distribution.
sentencepiece_model_path: str, path to the SentencePiece model file to use
for decoding. Must match the one used during training.
"""
if checkpoint_step == -1:
checkpoint_step = _get_latest_checkpoint_from_dir(self._model_dir)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
gin.bind_parameter("Bitransformer.decode.beam_size", beam_size)
gin.bind_parameter("Bitransformer.decode.temperature", temperature)
gin.bind_parameter("utils.get_variable_dtype.slice_dtype", "float32")
gin.bind_parameter("utils.get_variable_dtype.activation_dtype", "float32")
vocabulary = t5.data.SentencePieceVocabulary(sentencepiece_model_path)
model_ckpt = "model.ckpt-" + str(checkpoint_step)
export_dir = export_dir or self._model_dir
utils.export_model(
self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
self._sequence_length, batch_size=self.batch_size,
checkpoint_path=os.path.join(self._model_dir, model_ckpt))
def export(self, export_dir=None, checkpoint_step=-1, beam_size=1,
temperature=1.0, vocabulary=None):
"""Exports a TensorFlow SavedModel.
Args:
export_dir: str, a directory in which to export SavedModels. Will use
`model_dir` if unspecified.
checkpoint_step: int, checkpoint to export. If -1 (default), use the
latest checkpoint from the pretrained model directory.
beam_size: int, a number >= 1 specifying the number of beams to use for
beam search.
temperature: float, a value between 0 and 1 (must be 0 if beam_size > 1)
0.0 means argmax, 1.0 means sample according to predicted distribution.
vocabulary: vocabularies.Vocabulary object to use for tokenization, or
None to use the default SentencePieceVocabulary.
Returns:
The string path to the exported directory.
"""
if checkpoint_step == -1:
checkpoint_step = _get_latest_checkpoint_from_dir(self._model_dir)
with gin.unlock_config():
gin.parse_config_file(_operative_config_path(self._model_dir))
gin.bind_parameter("Bitransformer.decode.beam_size", beam_size)
gin.bind_parameter("Bitransformer.decode.temperature", temperature)
if vocabulary is None:
vocabulary = t5.data.get_default_vocabulary()
model_ckpt = "model.ckpt-" + str(checkpoint_step)
export_dir = export_dir or self._model_dir
return utils.export_model(
self.estimator(vocabulary, disable_tpu=True), export_dir, vocabulary,
self._sequence_length, batch_size=self.batch_size,
checkpoint_path=os.path.join(self._model_dir, model_ckpt))
def export(self,
export_dir=None,
checkpoint_step=-1,
beam_size=1,
temperature=1.0,
keep_top_k=-1,
vocabulary=None,
eval_with_score=False):
"""Exports a TensorFlow SavedModel.
Args:
export_dir: str, a directory in which to export SavedModels. Will use
`model_dir` if unspecified.
checkpoint_step: int, checkpoint to export. If -1 (default), use the
latest checkpoint from the pretrained model directory.
beam_size: int, a number >= 1 specifying the number of beams to use for
beam search.
temperature: float, a value between 0 and 1 (must be 0 if beam_size > 1)
0.0 means argmax, 1.0 means sample according to predicted distribution.
keep_top_k: integer, a value between 1 and the vocabulary size. When
sampling, only pick tokens that are in the k most likely.
vocabulary: vocabularies.Vocabulary object to use for tokenization, or
None to use the default SentencePieceVocabulary.
eval_with_score: If True, compute log-likelihood scores of targets.
If False, do inference to generate outputs.
Returns:
The string path to the exported directory.
"""
if checkpoint_step == -1:
checkpoint_step = utils.get_latest_checkpoint_from_dir(
self._model_dir)
_parse_operative_config(self._model_dir)
with gin.unlock_config():
gin.bind_parameter("Bitransformer.decode.beam_size", beam_size)
gin.bind_parameter("Bitransformer.decode.temperature", temperature)
gin.bind_parameter("Bitransformer.decode.sampling_keep_top_k",
keep_top_k)
if vocabulary is None:
vocabulary = utils.get_vocabulary()
model_ckpt = "model.ckpt-" + str(checkpoint_step)
export_dir = export_dir or self._model_dir
estimator = self.estimator(vocabulary,
disable_tpu=True,
score_in_predict_mode=eval_with_score)
return mtf_utils.export_model(estimator,
export_dir,
vocabulary,
self._sequence_length,
self._model_type,
batch_size=self.batch_size,
checkpoint_path=os.path.join(
self._model_dir, model_ckpt),
eval_with_score=eval_with_score)
