Add topk arg to return topk items and scores at inference step #678
Conversation
Documentation previewhttps://nvidia-merlin.github.io/Transformers4Rec/review/pr-678 |
| outputs[name] = task( | ||
| body_outputs, targets=targets, training=training, testing=testing, **kwargs | ||
| ) | ||
| if isinstance(task, NextItemPredictionTask): |
There was a problem hiding this comment.
If the NextItemPredictionTask is always something that is created by user code. An alternative to passing to the foward method could be accepting top_k in the constructor __init__ method. That way we wouldn't need to have this condition here
There was a problem hiding this comment.
Good point. adding it to the NextItemPredictionTask constructor might be easier (as I show below) but, the reason we did not add this to the constructor is bcs there is no turning back changing the topK value at the inf step, once the model is trained with top_k arg. So to make it flexible, we decided to add it in the forward method of NextItemPredictionTask and constructor of the Model class.
head = tr.Head(
body,
tr.NextItemPredictionTask(weight_tying=True,
metrics=metrics, top_k =20),
inputs=inputs,
)
@sararb wanna add something else?
| @@ -342,7 +344,11 @@ def forward(self, inputs: torch.Tensor, targets=None, training=False, testing=Fa | |||
| # Compute predictions probs | |||
| x, _ = self.pre(x) # type: ignore | |||
|
|
|||
| return x | |||
| if top_k == -1: | |||
There was a problem hiding this comment.
Does top_k need to be a numeric value? Would it work with with None as the default value, for example?
There was a problem hiding this comment.
I think -1 is a used convention but we can set it to None either. @sararb will setting it to None create any issue?
There was a problem hiding this comment.
You're right, it is just a convention. Setting it to None won't create any issue.
rnyak
changed the title
| @@ -342,7 +344,11 @@ def forward(self, inputs: torch.Tensor, targets=None, training=False, testing=Fa | |||
| # Compute predictions probs | |||
| x, _ = self.pre(x) # type: ignore | |||
|
|
|||
| return x | |||
| if top_k == -1: | |||
There was a problem hiding this comment.
You're right, it is just a convention. Setting it to None won't create any issue.
|
rerun tests |
1 similar comment
|
rerun tests |
This PR adds functionality for returning topk most relevant (with the highest scores) item ids from Triton IS, for NextItemPrediction task.
Current blocker:
The code designed to return top_k item ids (int 64 dtype), but model.output_schema returnsnext_itemas float32 dtype, which creates an error from Triton.Shall we change the code base in a way that model.output_schema matches with the expected output and output dtype from Triton? Or shall we return top_k item id scores, instead of item_ids?
Status update:
After modifying the model.output_schema, we can now return two outputs (item_scores, item_ids) from Triton.
Remaining tasks: