dataset.py (7406B)
1 from device import device 2 from torch.utils.data import Dataset 3 from torch_geometric.data import Data 4 from torch_geometric.utils.convert import from_scipy_sparse_matrix 5 from tqdm import trange 6 from transformers import AutoTokenizer, AutoModel 7 from typing import List, Dict 8 from utils import load_graph_from_path 9 import json 10 import networkx as nx 11 import torch 12 PLM = 'bert-base-uncased' 13 14 15 class TextGraphDataset(Dataset): 16 def __init__(self, data, args, head_emb=None, tail_emb=None, tokenizer=None): 17 self.special_token = "[SPECIAL_TOKEN]" 18 self.special_node = "SPECIAL_NODE]" 19 self.questions: List[str] = [] 20 self.question_paths: List[List[str]] = [] 21 self.paths: List[List[List[str], List[str]]] = [] 22 self.labels: List[int] = [] 23 self.sample_ids: List[str] = [] 24 self.graphs: List[Data] = [] 25 self.indicies: List[List[int]] = [] 26 self.unidirectional: bool = args.unidirectional 27 self.bidirectional: bool = args.bidirectional 28 self.model = AutoModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device) 29 node_embedding_size: int = self.model.config.hidden_size 30 31 if tokenizer is None: 32 self.tokenizer = AutoTokenizer.from_pretrained(PLM, use_fast=False) 33 if args.bidirectional: 34 assert args.unidirectional, "Unidirectional is set to false!" 35 self.tokenizer.add_tokens(self.special_token, special_tokens=True) 36 self.special_token_id = len(self.tokenizer) - 1 37 else: 38 self.tokenizer = tokenizer 39 40 self.texts: Dict[str, List[int]] = { 41 "input_ids": [], 42 "attention_mask": [] 43 } 44 45 if self.bidirectional: 46 self.model.resize_token_embeddings(len(self.tokenizer)) 47 48 if self.unidirectional: 49 self.special_node_indicies: List[List[int, int]] = [] 50 51 if head_emb is None and tail_emb is None: 52 self.head_emb: torch.Tensor = torch.randn(node_embedding_size) 53 self.tail_emb: torch.Tensor = torch.randn(node_embedding_size) 54 else: 55 self.head_emb: torch.Tensor = head_emb 56 self.tail_emb: torch.Tensor = tail_emb 57 58 with open("static/edge_type_dict.json", 'r') as f: 59 e2i = json.load(f) 60 61 for i in trange(len(data)): 62 graph_path: str = data[i]["graph"] 63 raw_graph: nx.MultiDiGraph = load_graph_from_path(graph_path) 64 question: str = data[i]["question"] 65 self.labels.append(data[i]["label"]) 66 self.sample_ids.append(i) 67 self.questions.append(question) 68 self.question_paths.append(data[i]["question_path"]) 69 self.paths.append(data[i]["og_path"]) 70 71 if self.bidirectional: 72 question += self.special_token 73 74 tokenized_question = self.tokenizer(question, padding="max_length", max_length=64) 75 question_input_ids = torch.tensor(tokenized_question["input_ids"]).view(1, -1).to(device) 76 question_attention_mask = torch.tensor(tokenized_question["attention_mask"]).view(1, -1).to(device) 77 self.texts["input_ids"].append(question_input_ids[0]) 78 self.texts["attention_mask"].append(question_attention_mask[0]) 79 80 head_string = data[i]["og_path"][0][0] 81 tail_string = data[i]["og_path"][-1][-1] 82 if args.GNN_only or args.static: 83 head = self.head_emb.tolist() 84 tail = self.tail_emb.tolist() 85 else: 86 head_tokenized = self.tokenizer(question, head_string, padding="max_length", max_length=64) 87 tail_tokenized = self.tokenizer(question, tail_string, padding="max_length", max_length=64) 88 head_input_ids = torch.tensor(head_tokenized["input_ids"]).view(1, -1).to(device) 89 tail_input_ids = torch.tensor(tail_tokenized["input_ids"]).view(1, -1).to(device) 90 head_attention_mask = torch.tensor(head_tokenized["attention_mask"]).view(1, -1).to(device) 91 tail_attention_mask = torch.tensor(tail_tokenized["attention_mask"]).view(1, -1).to(device) 92 with torch.no_grad(): 93 head_output = self.model(head_input_ids, head_attention_mask)[1].squeeze(0) 94 tail_output = self.model(tail_input_ids, tail_attention_mask)[1].squeeze(0) 95 head = head_output.tolist() 96 tail = tail_output.tolist() 97 98 if self.unidirectional or self.bidirectional: 99 with torch.no_grad(): 100 question_embedded = self.model(question_input_ids, question_attention_mask)[1].squeeze(0) 101 raw_graph.add_node(self.special_node) 102 for node in raw_graph.nodes(): 103 if node == self.special_node: 104 continue 105 raw_graph.add_edge(self.special_node, node, type="special", id="special") 106 raw_graph.add_edge(node, self.special_node, type="special_r", id="special_r") 107 108 edge_type_tensor = torch.tensor([e2i[e[-1]["type"]] for e in raw_graph.edges(data=True)], dtype=torch.long) 109 adj_matrix = nx.adjacency_matrix(raw_graph) 110 edge_index = from_scipy_sparse_matrix(adj_matrix) 111 node_embeddings = {n: torch.zeros(node_embedding_size).tolist() for i, n in enumerate(list(raw_graph.nodes()))} 112 node_embeddings[head_string] = head 113 node_embeddings[tail_string] = tail 114 node_keys = list(node_embeddings.keys()) 115 head_index = node_keys.index(head_string) 116 tail_index = node_keys.index(tail_string) 117 if self.unidirectional or self.bidirectional: 118 node_embeddings[self.special_node] = question_embedded 119 special_node_index = node_keys.index(self.special_node) 120 self.special_node_indicies.append(special_node_index) 121 nodes = torch.tensor(list(node_embeddings.values()), dtype=torch.float) 122 processed_graph = Data(x=nodes, edge_index=edge_index[0], edge_type=edge_type_tensor).to(device) 123 self.graphs.append(processed_graph) 124 self.indicies.append([head_index, tail_index]) 125 126 def __len__(self): 127 return len(self.labels) 128 129 def __getitem__(self, idx): 130 input_ids = self.texts["input_ids"][idx] 131 attention_mask = self.texts["attention_mask"][idx] 132 label = self.labels[idx] 133 134 if self.unidirectional or self.bidirectional: 135 special_node_index = self.special_node_indicies 136 else: 137 special_node_index = [0 for _ in self.indicies] 138 139 if self.bidirectional: 140 token = self.tokenizer(self.special_token, add_special_tokens=False)["input_ids"][0] 141 special_token_index = input_ids.tolist().index(token) 142 else: 143 special_token_index = [0 for _ in self.indicies] 144 145 return self.sample_ids[idx], torch.tensor(input_ids), torch.tensor(attention_mask), self.graphs[idx], self.indicies[idx], torch.tensor(special_node_index[idx]), torch.tensor(special_token_index), torch.tensor(label, dtype=torch.float32) 146 147 def decode(self, idx): 148 return self.tokenizer.decode(self.texts["input_ids"][idx], skip_special_tokens=False), self.labels[idx]