tfm.nlp.serving_modules.MaskedLM

View source on GitHub

The export module for the Bert Pretrain (MaskedLM) task.

Inherits From: ExportModule

tfm.nlp.serving_modules.MaskedLM(
    params, model: tf.keras.Model, inference_step=None
)

Args
params	A dataclass for parameters to the module.
model	A model instance which contains weights and forward computation.
inference_step	An optional callable to forward-pass the model. If not specified, it creates a parital function with model as an required kwarg.
preprocessor	An optional callable to preprocess the inputs.
postprocessor	An optional callable to postprocess the model outputs.

Child Classes

class Params

Methods

get_inference_signatures

View source

get_inference_signatures(
    function_keys: Dict[Text, Text]
)

Get defined function signatures.

serve

View source

@tf.function
serve(
    input_word_ids, input_mask, input_type_ids, masked_lm_positions
) -> Dict[str, tf.Tensor]

serve_examples

View source

@tf.function
serve_examples(
    inputs
) -> Dict[str, tf.Tensor]