Quotes Glove Model - Medium

Quotes_Glove_model.py

def get_batch(batch_size):
    ind = np.random.permutation(occs.size).tolist()
    i = 0
    for i in range(0, tot_pairs, batch_size):
        batch_ids = ind[i:i+batch_size]
        yield p1[batch_ids], p2[batch_ids], occs[batch_ids]

device = None
if torch.cuda.is_available():
    device = torch.device("cuda:0")
else:
    device = torch.device("cpu")

class Glove(nn.Module) :
    def __init__(self, vocab_len, num_dim,):
        super(Glove, self).__init__()
        self.ui = nn.Embedding(vocab_len, num_dim).to(device)
        self.uj = nn.Embedding(vocab_len, num_dim).to(device)
        self.bi = nn.Embedding(vocab_len, 1).to(device)
        self.bj = nn.Embedding(vocab_len, 1).to(device)

        self.ui.weight.data.uniform_(-1,1)
        self.uj.weight.data.uniform_(-1,1)
        self.bi.weight.data.zero_()
        self.bj.weight.data.zero_()
    
    def forward(self, i_vecs, j_vecs):
        i_s = self.ui(i_vecs)
        j_s = self.uj(j_vecs)
        b_i = torch.squeeze(self.bi(i_vecs))
        b_j = torch.squeeze(self.bj(j_vecs))
        return torch.sum(i_s * j_s, dim = 1) + b_i + b_j