tfr.extension.premade.TFRBertModelBuilder

Model builder for TFR-BERT models.

Inherits From: ModelBuilder, ModelBuilderWithMask, AbstractModelBuilder

View aliases

Main aliases

tfr.extension.premade.tfrbert_task.TFRBertModelBuilder

tfr.extension.premade.TFRBertModelBuilder(
    input_creator: Callable[[], Tuple[TensorDict, TensorDict]],
    preprocessor: Callable[[TensorDict, TensorDict, tf.Tensor], Tuple[TensorDict, TensorDict]],
    scorer: Callable[[TensorDict, TensorDict, tf.Tensor], Union[TensorLike, TensorDict]],
    mask_feature_name: str,
    name: Optional[str] = None
)

Args
input_creator	A callable or a class like InputCreator to implement create_inputs.
preprocessor	A callable or a class like Preprocessor to implement preprocess.
scorer	A callable or a class like Scorer to implement score.
mask_feature_name	name of 2D mask boolean feature.
name	(optional) name of the Model.

Methods

build

View source

build() -> tf.keras.Model

Builds a Keras Model for Ranking Pipeline.

Example usage:

model_builder = SimpleModelBuilder(
    {},
    {"example_feature_1": tf.io.FixedLenFeature(
        shape=(1,), dtype=tf.float32, default_value=0.0)},
    "list_mask", "model_builder")
model = model_builder.build()

Returns
A tf.keras.Model.

create_inputs

View source

create_inputs() -> Tuple[tfr.keras.model.TensorDict, tfr.keras.model.TensorDict, tf.Tensor]

See ModelBuilderWithMask.

preprocess

View source

preprocess(
    context_inputs: tfr.keras.model.TensorDict,
    example_inputs: tfr.keras.model.TensorDict,
    mask: tf.Tensor
) -> Tuple[tfr.keras.model.TensorDict, tfr.keras.model.TensorDict]

See ModelBuilderWithMask.

score

View source

score(
    context_features: tfr.keras.model.TensorDict,
    example_features: tfr.keras.model.TensorDict,
    mask: tf.Tensor
) -> Union[TensorLike, TensorDict]

See ModelBuilderWithMask.