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)