from torch import nn
import torch

class OurModel(nn.Module):
    def __init__(self, num_labels, emb_dim=10):
        super(OurModel, self).__init__()
        self.emb = nn.Embedding(num_labels, emb_dim)
        self.user_embs = nn.Embedding(1, emb_dim)

    def forward(self, pos, neg):
        batch_size = neg.size(0)
        one_user_vector = self.user_embs(torch.zeros(1).long())
        repeated_user_vector = one_user_vector.repeat((batch_size, 1)).view(batch_size, -1, 1)
        pos_res = torch.bmm(self.emb(pos), repeated_user_vector).squeeze(2)
        neg_res = torch.bmm(self.emb(neg), repeated_user_vector).squeeze(2)

        return pos_res, neg_res

num_labels = 100
model = OurModel(num_labels)

pos_labels = torch.randint(high=num_labels, size=(3,1)) # our five labels
neg_labels = torch.randint(high=num_labels, size=(3,2)) # a few random negatives per positive

pos_res, neg_res = model(pos_labels, neg_labels)
print('Positive Labels:', pos_labels)
print('Negative Labels:', neg_labels)
print('Model positive scores:', pos_res)
print('Model negative scores:', neg_res)
loss = warp_loss(pos_res, neg_res, num_labels=num_labels, device=torch.device('cpu'))
print('Loss:', loss)
loss.backward()

print('We can also see that the gradient is only active for 2x the number of positive labels:', (model.emb.weight.grad.sum(1) != 0).sum().item())
print('Meaning we correctly discard the gradients for all other than the offending negative label.')