Infer step size for Embeddings

[KuuCi]

Theres two cases:

1. For in batch negatives, total number of negatives is batch size - 1. And pos_step_size is 2.
2. For hard negatives, each sample has a list of negatives that can be as long as you like. In this case you have to change the pos_step_size.

We can remove setting step_size and autoinfer without needing the pos_step_size. We will need to look at the first batch to determine how many hard negatives. If there are hard negatives, we will use that - 1 as the step size. Otherwise, we will defer to batch negatives, setting step_size = 2

Testing:

embedding-ft-no-pos-zXOtvz (off of this branch, removed pos_step_size keyword):

embedding-ft-pos-kl2fPy (off foundry main pos_step_size = 2):

embedding-ft-pos-neg-gMJSGx (off foundry main pos_step_size = 21)

embedding-ft-no-pos-neg-3GoEms (off of this branch, removed pos_step_size keyword):

[jacobfulano]

Small comments

[mrdrozdov]

I really like this idea. I do have a couple soft suggestions:

1. Rather than self._first_batch_seen can you simply set self.step_size = None during init? Then check if self.step_size is None instead of using first_batch_seen.

2. It could be more tidy to just make step_size part of the batch. I think you could do something like batch["step_size"] = step_size, so that it is accessible in compute_score? That being said, I'm not familiar enough to foundry to know if this causes complications, and you might need to do something like batch["step_size"] = torch.tensor([step_size] * batch_size, dtype=torch.long). If you go this route, then you can ignore (1).

Anyway, the current approach is fine too, but wanted to share just in case.

[mrdrozdov]

This looks mostly good to me!

A few requests for integration tests:

• Confirm expected behavior when training data has no hard negatives.
• Confirm when using hard negatives = 1.
• Confirm when using hard negatives = 3 (or some number besides 1).

Maybe you've already done this? In which case I will switch to approve after you address my minor comments.

[KuuCi]

Tested with the following runs:

0 negatives: embedding-ft-no-pos-neg-0-7jyk87

1 negative: embedding-ft-no-pos-neg-1-zAeAzg

20 negatives: embedding-ft-no-pos-neg-20-EicOP8

All looks good

[mrdrozdov]