def forward(self, user: TensorType, doc: TensorType) -> TensorType:
"""Evaluate the user-doc Q model
Args:
user: User embedding of shape (batch_size, user embedding size).
Note that `self.embedding_size` is the sum of both user- and
doc-embedding size.
doc: Doc embeddings of shape (batch_size, num_docs, doc embedding size).
Note that `self.embedding_size` is the sum of both user- and
doc-embedding size.
Returns:
The q_values per document of shape (batch_size, num_docs + 1). +1 due to
also having a Q-value for the non-interaction (no click/no doc).
"""
batch_size, num_docs, embedding_size = doc.shape
doc_flat = doc.view((batch_size * num_docs, embedding_size))
# Concat everything.
# No user features.
if user.shape[-1] == 0:
x = doc_flat
# User features, repeat user embeddings n times (n=num docs).
else:
user_repeated = user.repeat(num_docs, 1)
x = torch.cat([user_repeated, doc_flat], dim=1)
x = self.layers(x)
# Similar to Google's SlateQ implementation in RecSim, we force the
# Q-values to zeros if there are no clicks.
# See https://arxiv.org/abs/1905.12767 for details.
x_no_click = torch.zeros((batch_size, 1), device=x.device)
return torch.cat([x.view((batch_size, num_docs)), x_no_click], dim=1)
def forward(self, user: TensorType, doc: TensorType) -> TensorType:
"""Evaluate the user-doc Q model
Args:
user (TensorType): User embedding of shape (batch_size,
embedding_size).
doc (TensorType): Doc embeddings of shape (batch_size, num_docs,
embedding_size).
Returns:
score (TensorType): q_values of shape (batch_size, num_docs + 1).
"""
batch_size, num_docs, embedding_size = doc.shape
doc_flat = doc.view((batch_size * num_docs, embedding_size))
user_repeated = user.repeat(num_docs, 1)
x = torch.cat([user_repeated, doc_flat], dim=1)
x = self.layers(x)
# Similar to Google's SlateQ implementation in RecSim, we force the
# Q-values to zeros if there are no clicks.
x_no_click = torch.zeros((batch_size, 1), device=x.device)
return torch.cat([x.view((batch_size, num_docs)), x_no_click], dim=1)
