How can I use sample_weights for retrieval task?

[modic23]

Hi,

I'm trying to build a retrieval model. I have positive samples and some weights corresponding to each of these sample which helps to denote that for the given sample how much weight it holds. I want to use these weights so that I have better recommendations from the retrieval model. I'm assuming that sample_weights param can help me to do this.

Let me know if I'm thinking in the right way and if yes, Is this the right way to do this?

self.task = tfrs.tasks.Retrieval(metrics=metrics)
def compute_loss(self, features, training=False):
q_emb = ...
p_emb = ...
return self.task(q_emb, p_emb, compute_metrics=not training, sample_weight=features['weight'])