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)