YuWang0103/LGGM-Text2Graph

1import torch2 3 4class DistributionNodes:5    def __init__(self, histogram):6        """ Compute the distribution of the number of nodes in the dataset, and sample from this distribution.7            historgram: dict. The keys are num_nodes, the values are counts8        """9 10        if type(histogram) == dict:11            max_n_nodes = max(histogram.keys())12            prob = torch.zeros(max_n_nodes + 1)13            for num_nodes, count in histogram.items():14                prob[num_nodes] = count15        else:16            prob = histogram17 18        self.prob = prob / prob.sum()19        self.m = torch.distributions.Categorical(prob)20 21    def sample_n(self, n_samples, device):22        idx = self.m.sample((n_samples,))23        return idx.to(device)24 25    def log_prob(self, batch_n_nodes):26        assert len(batch_n_nodes.size()) == 127        p = self.prob.to(batch_n_nodes.device)28 29        mask = batch_n_nodes >= p.shape[0]30        batch_n_nodes[mask] = p.shape[0] - 131 32        probas = p[batch_n_nodes]33 34        probas[mask] = 035        log_p = torch.log(probas + 1e-30)36        37        return log_p