YuWang0103/LGGM-Text2Graph

1import torch2import torch.nn as nn3import torch.nn.functional as F4from tqdm import tqdm5 6from models.transformer_model import GraphTransformer7from diffusion.noise_schedule import DiscreteUniformTransition, PredefinedNoiseScheduleDiscrete8from diffusion import diffusion_utils9import utils10import networkx as nx11from sentence_transformers import SentenceTransformer12import pytorch_lightning as pl13from transformers import BertTokenizer, BertForSequenceClassification14 15 16class LGGMText2Graph_Demo(pl.LightningModule):17    def __init__(self, cfg, input_dims, output_dims, cond_dims, cond_emb, \18                 nodes_dist, node_types, edge_types, extra_features, data_loaders):19        super().__init__()20 21        nodes_dist = nodes_dist22 23        self.cfg = cfg24        self.T = cfg.model.diffusion_steps25 26        self.Xdim = input_dims['X']27        self.Edim = input_dims['E']28        self.ydim = input_dims['y']29        self.Xdim_output = output_dims['X']30        self.Edim_output = output_dims['E']31        self.ydim_output = output_dims['y']32        self.node_dist = nodes_dist33 34 35        self.extra_features = extra_features36 37        self.model = GraphTransformer(n_layers=cfg.model.n_layers,38                                      input_dims=input_dims,39                                      hidden_mlp_dims=cfg.model.hidden_mlp_dims,40                                      hidden_dims=cfg.model.hidden_dims,41                                      output_dims=output_dims,42                                      cond_dims = cond_dims,43                                      act_fn_in=nn.ReLU(),44                                      act_fn_out=nn.ReLU()).to(self.device)45        46 47        self.noise_schedule = PredefinedNoiseScheduleDiscrete(cfg.model.diffusion_noise_schedule,48                                                              timesteps=cfg.model.diffusion_steps).to(self.device)49 50        self.transition_model = DiscreteUniformTransition(x_classes=self.Xdim_output, e_classes=self.Edim_output,51                                                            y_classes=self.ydim_output)52        x_limit = torch.ones(self.Xdim_output) / self.Xdim_output53        e_limit = torch.ones(self.Edim_output) / self.Edim_output54        y_limit = torch.ones(self.ydim_output) / self.ydim_output55        56        self.limit_dist = utils.PlaceHolder(X=x_limit, E=e_limit, y=y_limit)57        58 59    def generate_basic(self, text, num_nodes) -> None:60        print(num_nodes)61        prompt_emb = torch.tensor(self.text_encoder.encode([text])).to(self.device)62        samples = self.sample_batch(5, cond_emb = prompt_emb, num_nodes = num_nodes)63 64        nx_graphs = []65        for graph in samples:66            node_types, edge_types = graph67            A = edge_types.bool().cpu().numpy()68 69            nx_graph = nx.from_numpy_array(A)70            nx_graphs.append(nx_graph)71        72        return nx_graphs73 74    def generate_pretrained(self, text, num_nodes) -> None:75        encoded_input = self.tokenizer(text, return_tensors='pt', padding=True, truncation=True, max_length=512)76        encoded_input = {key: val.to(self.text_encoder.device) for key, val in encoded_input.items()}77 78        # Get the model output79        with torch.no_grad():80            prompt_emb = self.text_encoder(**encoded_input).hidden_states[-1][:, 0]81 82        samples = self.sample_batch(3, cond_emb = prompt_emb.to(self.device), num_nodes = num_nodes)83 84        nx_graphs = []85        for graph in samples:86            node_types, edge_types = graph87            A = edge_types.bool().cpu().numpy()88 89            nx_graph = nx.from_numpy_array(A)90            nx_graphs.append(nx_graph)91        92        return nx_graphs93 94    def init_prompt_encoder_basic(self):95        self.text_encoder = SentenceTransformer("all-MiniLM-L6-v2")96    97    def init_prompt_encoder_pretrained(self):98        model_name = f"./checkpoint-900"  # or "bert-base-uncased" if starting from the base model99        self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')100        self.text_encoder = BertForSequenceClassification.from_pretrained(model_name, num_labels=8, output_hidden_states=True, device_map = 'cpu')101 102 103    @torch.no_grad()104    def sample_batch(self, batch_size: int,  cond_emb = None, num_nodes = None):105        """106        :param batch_id: int107        :param batch_size: int108        :param num_nodes: int, <int>tensor (batch_size) (optional) for specifying number of nodes109        :param save_final: int: number of predictions to save to file110        :param keep_chain: int: number of chains to save to file111        :param keep_chain_steps: number of timesteps to save for each chain112        :return: molecule_list. Each element of this list is a tuple (atom_types, charges, positions)113        """114        if num_nodes is None:115            n_nodes = self.node_dist.sample_n(batch_size, self.device)116        elif type(num_nodes) == int:117            n_nodes = num_nodes * torch.ones(batch_size, device=self.device, dtype=torch.int)118        119        n_max = torch.max(n_nodes).item()120        # Build the masks121        arange = torch.arange(n_max, device=self.device).unsqueeze(0).expand(batch_size, -1)122        node_mask = arange < n_nodes.unsqueeze(1)123        # Sample noise  -- z has size (n_samples, n_nodes, n_features)124        125        z_T = diffusion_utils.sample_discrete_feature_noise(limit_dist=self.limit_dist, node_mask=node_mask, transition=self.cfg.model.transition)126        X, E, y = z_T.X, z_T.E, z_T.y127 128 129        # Iteratively sample p(z_s | z_t) for t = 1, ..., T, with s = t - 1.130        for s_int in tqdm(reversed(range(0, self.T))):131            s_array = s_int * torch.ones((batch_size, 1)).type_as(y)132            t_array = s_array + 1133            s_norm = s_array / self.T134            t_norm = t_array / self.T135 136            # Sample z_s137            sampled_s = self.sample_p_zs_given_zt(s_norm, t_norm, X, E, y, node_mask, cond_emb)138            X, E, y = sampled_s.X, sampled_s.E, sampled_s.y139 140        # Sample141        sampled_s = sampled_s.mask(node_mask, collapse=True)142        X, E, y = sampled_s.X, sampled_s.E, sampled_s.y143 144 145        graph_list = []146        for i in range(batch_size):147            n = n_nodes[i]148            node_types = X[i, :n].cpu()149            edge_types = E[i, :n, :n].cpu()150            graph_list.append([node_types, edge_types])151 152        return graph_list153 154    def sample_p_zs_given_zt(self, s, t, X_t, E_t, y_t, node_mask, cond_emb):155        """Samples from zs ~ p(zs | zt). Only used during sampling.156           if last_step, return the graph prediction as well"""157        bs, n, dxs = X_t.shape158        beta_t = self.noise_schedule(t_normalized=t)  # (bs, 1)159        alpha_s_bar = self.noise_schedule.get_alpha_bar(t_normalized=s)160        alpha_t_bar = self.noise_schedule.get_alpha_bar(t_normalized=t)161 162        163        # Retrieve transitions matrix164        Qtb = self.transition_model.get_Qt_bar(alpha_t_bar, self.device)165        Qsb = self.transition_model.get_Qt_bar(alpha_s_bar, self.device)166        Qt = self.transition_model.get_Qt(beta_t, self.device)167 168        noisy_data = {'X_t': X_t, 'E_t': E_t, 'y_t': y_t, 't': t, 'node_mask': node_mask, 'cond_emb': cond_emb.repeat(X_t.shape[0], 1)}169        extra_data = self.compute_extra_data(noisy_data)170        pred = self.forward(noisy_data, extra_data, node_mask)171 172        # Normalize predictions173        pred_X = F.softmax(pred.X, dim=-1)               # bs, n, d0174        pred_E = F.softmax(pred.E, dim=-1)               # bs, n, n, d0175 176        p_s_and_t_given_0_X = diffusion_utils.compute_batched_over0_posterior_distribution(X_t=X_t,177                                                                                           Qt=Qt.X,178                                                                                           Qsb=Qsb.X,179                                                                                           Qtb=Qtb.X)180 181        p_s_and_t_given_0_E = diffusion_utils.compute_batched_over0_posterior_distribution(X_t=E_t,182                                                                                           Qt=Qt.E,183                                                                                           Qsb=Qsb.E,184                                                                                           Qtb=Qtb.E)185        # Dim of these two tensors: bs, N, d0, d_t-1186        weighted_X = pred_X.unsqueeze(-1) * p_s_and_t_given_0_X         # bs, n, d0, d_t-1187        unnormalized_prob_X = weighted_X.sum(dim=2)                     # bs, n, d_t-1188        unnormalized_prob_X[torch.sum(unnormalized_prob_X, dim=-1) == 0] = 1e-5189        prob_X = unnormalized_prob_X / torch.sum(unnormalized_prob_X, dim=-1, keepdim=True)  # bs, n, d_t-1190 191        pred_E = pred_E.reshape((bs, -1, pred_E.shape[-1]))192        weighted_E = pred_E.unsqueeze(-1) * p_s_and_t_given_0_E        # bs, N, d0, d_t-1193        unnormalized_prob_E = weighted_E.sum(dim=-2)194        unnormalized_prob_E[torch.sum(unnormalized_prob_E, dim=-1) == 0] = 1e-5195        prob_E = unnormalized_prob_E / torch.sum(unnormalized_prob_E, dim=-1, keepdim=True)196        prob_E = prob_E.reshape(bs, n, n, pred_E.shape[-1])197 198        assert ((prob_X.sum(dim=-1) - 1).abs() < 1e-4).all()199        assert ((prob_E.sum(dim=-1) - 1).abs() < 1e-4).all()200 201 202        sampled_s = diffusion_utils.sample_discrete_features(prob_X, prob_E, node_mask=node_mask)203 204 205        X_s = F.one_hot(sampled_s.X, num_classes=self.Xdim_output).float()206        E_s = F.one_hot(sampled_s.E, num_classes=self.Edim_output).float()207 208        assert (E_s == torch.transpose(E_s, 1, 2)).all()209        assert (X_t.shape == X_s.shape) and (E_t.shape == E_s.shape)210 211        out_one_hot = utils.PlaceHolder(X=X_s, E=E_s, y=torch.zeros(y_t.shape[0], 0))212 213        return out_one_hot.mask(node_mask).type_as(y_t)214 215    def compute_extra_data(self, noisy_data):216        """ At every training step (after adding noise) and step in sampling, compute extra information and append to217            the network input. """218 219        extra_features = self.extra_features(noisy_data)220 221        # print(extra_features.X.shape, extra_features.E.shape, extra_features.y.shape)222        extra_X = extra_features.X223        extra_E = extra_features.E224        extra_y = extra_features.y225 226        t = noisy_data['t']227        extra_y = torch.cat((extra_y, t), dim=1)228 229        return utils.PlaceHolder(X=extra_X, E=extra_E, y=extra_y)230 231    def forward(self, noisy_data, extra_data, node_mask):232        # print(noisy_data['cond_emb'].sum())233        B = noisy_data['cond_emb'].unsqueeze(1).unsqueeze(2).expand(-1, noisy_data['X_t'].shape[1], noisy_data['X_t'].shape[1], -1).to(self.device)234        A = noisy_data['cond_emb'].unsqueeze(1).expand(-1, noisy_data['X_t'].shape[1], -1).to(self.device)235 236        X = torch.cat((noisy_data['X_t'], extra_data.X, A), dim=2).float()237        E = torch.cat((noisy_data['E_t'], extra_data.E, B), dim=3).float()238        y = torch.hstack((noisy_data['y_t'], extra_data.y)).float()239 240        return self.model(X, E, y, node_mask)241