"""
Documentation
"""
# Python Modules
from functools import partial
from typing import Any, Dict
# 3rd Party Modules
import tensorflow as tf
import tensorflow.keras.initializers as tf_initializers
import tensorflow.keras.layers as tf_layers
# Project Modules
from deletor.models.utils import normalize_dropout
[docs]class ModelParameter(object):
# (optional int) Size of the output layer (i.e., the maximum number of documents
# a query can have). If None, use all documents in a query. Note each,
# batch may have a different shape in this case.
LIST_SIZE = 'list_size'
"""
(**Optional**) The maximum number of documents to include for each query or ```None```
for all of them.
"""
# (List[int]) The number of units in each layer. This will also determine
# the architecture of the network (i.e., number of layers).
N_UNITS = 'n_units'
"""
(**Required**) The number of hidden units for each layer.
"""
# (optional float or List[float]) The dropout rate.
DROPOUT_RATE = 'dropout_rate'
"""
(**Optional**) The dropout rate for each layer. This can be a single number
in which case the same dropout will be applied for each layer. Or it
can be a list of the same size as ```N_UNITS```.
"""
# (optional int) A random seed to use for any stochastic operations.
RANDOM_SEED = 'random_seed'
"""
(**Optional**) A random seed to use for any stochastic operations.
"""
[docs]class SimpleScoringNetwork(tf.keras.Model):
"""
A baseline model that uses a fully connected dense network, with one output,
for ranking.
"""
default_share_weights = False
default_dropout_rate = 0.0
def __init__(self, params: Dict[str, Any], **kwargs):
"""
:param params: See :class:`.ModelParameter` for the valid parameters.
:param kwargs:
"""
super().__init__(**kwargs)
# Model parameters
list_size = params[ModelParameter.LIST_SIZE]
n_units = params[ModelParameter.N_UNITS]
n_layers = len(n_units)
dropout_rate = normalize_dropout(params.get(ModelParameter.DROPOUT_RATE), n_units)
random_seed = params.get(ModelParameter.RANDOM_SEED)
self.list_size = None if list_size is None else tf.constant(list_size, tf.int32)
self.random_seed = None if random_seed is None else tf.constant(random_seed, tf.int64)
self.n_layers = n_layers
self.weight_initializer = weight_initializer = tf_initializers.he_uniform(seed=random_seed)
# Model layers
dense_layer = partial(tf_layers.Dense, kernel_initializer=weight_initializer)
self.dense = [dense_layer(n) for n in n_units]
self.activation = [tf_layers.PReLU() for _ in range(n_layers)]
self.batch_norm = [tf_layers.BatchNormalization() for _ in range(n_layers)]
# self.has_dropout = tf.constant(any([r for r in dropout_rate]) > 0, tf.bool)
self.has_dropout = any([r > 0 for r in dropout_rate])
if self.has_dropout is True:
self.dropout = [tf_layers.Dropout(r) for r in dropout_rate]
self.scoring_layer = dense_layer(1)
# noinspection DuplicatedCode
[docs] def call(self, inputs, training: bool = True, **kwargs):
"""
:param inputs:
:param training:
:param kwargs:
:return:
"""
# Should be shape (batch_size, list_size, n_features)
# dense_features = inputs['sequential']['dense']
dense_features = inputs['sequence_dense']
batch_size = tf.shape(dense_features)[0]
list_size = tf.shape(dense_features)[1]
n_features = tf.shape(dense_features)[2]
# Reshape the input. Each document will be scored separately, but
# the loss will be computed over the entire list.
lyr = tf.reshape(dense_features, (-1, n_features))
for i in range(self.n_layers):
lyr = self.dense[i](lyr)
lyr = self.batch_norm[i](lyr, training=training)
lyr = self.activation[i](lyr)
if self.has_dropout is True:
lyr = self.dropout[i](lyr, training=training)
scores = self.scoring_layer(lyr)
scores = tf.reshape(scores, (batch_size, list_size))
return scores