def __getitem__()

in hype_kg/codes/dataloader.py [0:0]

• 23 lines of code
• 3 McCabe index (conditional complexity)

    def __getitem__(self, idx):
        query = self.triples[idx][:-2]
        subsampling_weight = self.count[query]
        subsampling_weight = torch.sqrt(1 / torch.Tensor([subsampling_weight]))
        negative_sample_list = []
        negative_sample_size = 0
        while negative_sample_size < self.negative_sample_size:
            negative_sample = np.random.randint(self.nentity, size=self.negative_sample_size*2)
            mask = np.in1d(
                negative_sample, 
                self.true_tail[query], 
                assume_unique=True, 
                invert=True
            )
            negative_sample = negative_sample[mask]
            negative_sample_list.append(negative_sample)
            negative_sample_size += negative_sample.size
        negative_sample = np.concatenate(negative_sample_list)[:self.negative_sample_size]
        negative_sample = torch.from_numpy(negative_sample)
        flat_query = np.array([[qi[0], qi[1][0]] for qi in query]).flatten()
        tail = np.random.choice(list(self.true_tail[query]))
        positive_sample = torch.LongTensor(list(flat_query)+[tail])
        return positive_sample, negative_sample, subsampling_weight, self.mode