commit c196be5171521e5c214fd029210e83eba57f5735
parent acd356df74078d67833a88ad4d6c513eeb9f8559
Author: Sondre Wold <sondrewold1@gmail.com>
Date: Thu, 30 May 2024 15:18:17 +0200
Add run files
Diffstat:
| A | bidirectional.py | | | 834 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | dataset.py | | | 148 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | device.py | | | 2 | ++ |
| A | model.py | | | 116 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | train.py | | | 221 | +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++ |
| A | utils.py | | | 28 | ++++++++++++++++++++++++++++ |
6 files changed, 1349 insertions(+), 0 deletions(-)
diff --git a/bidirectional.py b/bidirectional.py
@@ -0,0 +1,834 @@
+### All BERT related code copied from the Huggingface Transformers library.
+import torch
+import math
+from typing import List, Optional, Tuple, Union
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from transformers.models.bert.configuration_bert import BertConfig
+from transformers.utils import (
+ ModelOutput,
+ add_code_sample_docstrings,
+ add_start_docstrings,
+ add_start_docstrings_to_model_forward,
+ logging,
+ replace_return_docstrings,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.pytorch_utils import apply_chunking_to_forward, find_pruneable_heads_and_indices, prune_linear_layer
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import (
+ BaseModelOutputWithPastAndCrossAttentions,
+ BaseModelOutputWithPoolingAndCrossAttentions,
+ CausalLMOutputWithCrossAttentions,
+ MaskedLMOutput,
+ MultipleChoiceModelOutput,
+ NextSentencePredictorOutput,
+ QuestionAnsweringModelOutput,
+ SequenceClassifierOutput,
+ TokenClassifierOutput,
+)
+def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
+ """Load tf checkpoints in a pytorch model."""
+ try:
+ import re
+
+ import numpy as np
+ import tensorflow as tf
+ except ImportError:
+ logger.error(
+ "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+ "https://www.tensorflow.org/install/ for installation instructions."
+ )
+ raise
+ tf_path = os.path.abspath(tf_checkpoint_path)
+ logger.info(f"Converting TensorFlow checkpoint from {tf_path}")
+ # Load weights from TF model
+ init_vars = tf.train.list_variables(tf_path)
+ names = []
+ arrays = []
+ for name, shape in init_vars:
+ logger.info(f"Loading TF weight {name} with shape {shape}")
+ array = tf.train.load_variable(tf_path, name)
+ names.append(name)
+ arrays.append(array)
+
+ for name, array in zip(names, arrays):
+ name = name.split("/")
+ # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+ # which are not required for using pretrained model
+ if any(
+ n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
+ for n in name
+ ):
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ pointer = model
+ for m_name in name:
+ if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
+ scope_names = re.split(r"_(\d+)", m_name)
+ else:
+ scope_names = [m_name]
+ if scope_names[0] == "kernel" or scope_names[0] == "gamma":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
+ pointer = getattr(pointer, "bias")
+ elif scope_names[0] == "output_weights":
+ pointer = getattr(pointer, "weight")
+ elif scope_names[0] == "squad":
+ pointer = getattr(pointer, "classifier")
+ else:
+ try:
+ pointer = getattr(pointer, scope_names[0])
+ except AttributeError:
+ logger.info(f"Skipping {'/'.join(name)}")
+ continue
+ if len(scope_names) >= 2:
+ num = int(scope_names[1])
+ pointer = pointer[num]
+ if m_name[-11:] == "_embeddings":
+ pointer = getattr(pointer, "weight")
+ elif m_name == "kernel":
+ array = np.transpose(array)
+ try:
+ if pointer.shape != array.shape:
+ raise ValueError(f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched")
+ except ValueError as e:
+ e.args += (pointer.shape, array.shape)
+ raise
+ logger.info(f"Initialize PyTorch weight {name}")
+ pointer.data = torch.from_numpy(array)
+ return model
+
+class BertPreTrainedModel(PreTrainedModel):
+ """
+ An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+ models.
+ """
+
+ config_class = BertConfig
+ load_tf_weights = load_tf_weights_in_bert
+ base_model_prefix = "bert"
+ supports_gradient_checkpointing = True
+
+ def _init_weights(self, module):
+ """Initialize the weights"""
+ if isinstance(module, nn.Linear):
+ # Slightly different from the TF version which uses truncated_normal for initialization
+ # cf https://github.com/pytorch/pytorch/pull/5617
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+ elif isinstance(module, nn.LayerNorm):
+ module.bias.data.zero_()
+ module.weight.data.fill_(1.0)
+
+
+class BertModel(BertPreTrainedModel):
+ """
+
+ The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+ cross-attention is added between the self-attention layers, following the architecture described in [Attention is
+ all you need](https://arxiv.org/abs/1706.03762) by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+ Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+
+ To behave as an decoder the model needs to be initialized with the `is_decoder` argument of the configuration set
+ to `True`. To be used in a Seq2Seq model, the model needs to initialized with both `is_decoder` argument and
+ `add_cross_attention` set to `True`; an `encoder_hidden_states` is then expected as an input to the forward pass.
+ """
+
+ def __init__(self, config, add_pooling_layer=True):
+ super().__init__(config)
+ self.config = config
+
+ self.embeddings = BertEmbeddings(config)
+ self.encoder = BertEncoder(config)
+
+ self.pooler = BertPooler(config) if add_pooling_layer else None
+
+ # Initialize weights and apply final processing
+ self.post_init()
+
+ def get_input_embeddings(self):
+ return self.embeddings.word_embeddings
+
+ def set_input_embeddings(self, value):
+ self.embeddings.word_embeddings = value
+
+ def _prune_heads(self, heads_to_prune):
+ """
+ Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+ class PreTrainedModel
+ """
+ for layer, heads in heads_to_prune.items():
+ self.encoder.layer[layer].attention.prune_heads(heads)
+
+ def forward(
+ self,
+ input_ids: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ token_type_ids: Optional[torch.Tensor] = None,
+ position_ids: Optional[torch.Tensor] = None,
+ head_mask: Optional[torch.Tensor] = None,
+ inputs_embeds: Optional[torch.Tensor] = None,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ encoder_attention_mask: Optional[torch.Tensor] = None,
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = None,
+ output_hidden_states: Optional[bool] = None,
+ return_dict: Optional[bool] = None,
+ special_node_embedding: Optional[torch.Tensor] = None,
+ spt: Optional[torch.Tensor] = None,
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPoolingAndCrossAttentions]:
+ r"""
+ encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+ Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention if
+ the model is configured as a decoder.
+ encoder_attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+ Mask to avoid performing attention on the padding token indices of the encoder input. This mask is used in
+ the cross-attention if the model is configured as a decoder. Mask values selected in `[0, 1]`:
+
+ - 1 for tokens that are **not masked**,
+ - 0 for tokens that are **masked**.
+ past_key_values (`tuple(tuple(torch.FloatTensor))` of length `config.n_layers` with each tuple having 4 tensors of shape `(batch_size, num_heads, sequence_length - 1, embed_size_per_head)`):
+ Contains precomputed key and value hidden states of the attention blocks. Can be used to speed up decoding.
+
+ If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that
+ don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all
+ `decoder_input_ids` of shape `(batch_size, sequence_length)`.
+ use_cache (`bool`, *optional*):
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+ `past_key_values`).
+ """
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if self.config.is_decoder:
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
+ else:
+ use_cache = False
+
+ if input_ids is not None and inputs_embeds is not None:
+ raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+ elif input_ids is not None:
+ self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
+ input_shape = input_ids.size()
+ elif inputs_embeds is not None:
+ input_shape = inputs_embeds.size()[:-1]
+ else:
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+ batch_size, seq_length = input_shape
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+ # past_key_values_length
+ past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0
+
+ if attention_mask is None:
+ attention_mask = torch.ones(((batch_size, seq_length + past_key_values_length)), device=device)
+
+ if token_type_ids is None:
+ if hasattr(self.embeddings, "token_type_ids"):
+ buffered_token_type_ids = self.embeddings.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(batch_size, seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+ # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+ # ourselves in which case we just need to make it broadcastable to all heads.
+ extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape)
+
+ # If a 2D or 3D attention mask is provided for the cross-attention
+ # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+ if self.config.is_decoder and encoder_hidden_states is not None:
+ encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+ encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+ if encoder_attention_mask is None:
+ encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+ encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+ else:
+ encoder_extended_attention_mask = None
+
+ # Prepare head mask if needed
+ # 1.0 in head_mask indicate we keep the head
+ # attention_probs has shape bsz x n_heads x N x N
+ # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+ # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+ head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+ embedding_output = self.embeddings(
+ input_ids=input_ids,
+ position_ids=position_ids,
+ token_type_ids=token_type_ids,
+ inputs_embeds=inputs_embeds,
+ past_key_values_length=past_key_values_length,
+ special_node_embedding=special_node_embedding,
+ spt=spt
+ )
+ encoder_outputs = self.encoder(
+ embedding_output,
+ attention_mask=extended_attention_mask,
+ head_mask=head_mask,
+ encoder_hidden_states=encoder_hidden_states,
+ encoder_attention_mask=encoder_extended_attention_mask,
+ past_key_values=past_key_values,
+ use_cache=use_cache,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+ sequence_output = encoder_outputs[0]
+ pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+ if not return_dict:
+ return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+ return BaseModelOutputWithPoolingAndCrossAttentions(
+ last_hidden_state=sequence_output,
+ pooler_output=pooled_output,
+ past_key_values=encoder_outputs.past_key_values,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ cross_attentions=encoder_outputs.cross_attentions,
+ )
+
+class BertEmbeddings(nn.Module):
+ """Construct the embeddings from word, position and token_type embeddings."""
+
+ def __init__(self, config):
+ super().__init__()
+ self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+ self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+ self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+ # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+ # any TensorFlow checkpoint file
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+ # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+ self.position_embedding_type = getattr(config, "position_embedding_type", "absolute")
+ self.register_buffer(
+ "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)), persistent=False
+ )
+ self.register_buffer(
+ "token_type_ids", torch.zeros(self.position_ids.size(), dtype=torch.long), persistent=False
+ )
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ token_type_ids: Optional[torch.LongTensor] = None,
+ position_ids: Optional[torch.LongTensor] = None,
+ inputs_embeds: Optional[torch.FloatTensor] = None,
+ past_key_values_length: int = 0,
+ special_node_embedding: Optional[torch.LongTensor] = None,
+ spt: Optional[torch.LongTensor] = None,
+ ) -> torch.Tensor:
+ if input_ids is not None:
+ input_shape = input_ids.size()
+ else:
+ input_shape = inputs_embeds.size()[:-1]
+
+ seq_length = input_shape[1]
+
+ if position_ids is None:
+ position_ids = self.position_ids[:, past_key_values_length : seq_length + past_key_values_length]
+
+ # Setting the token_type_ids to the registered buffer in constructor where it is all zeros, which usually occurs
+ # when its auto-generated, registered buffer helps users when tracing the model without passing token_type_ids, solves
+ # issue #5664
+ if token_type_ids is None:
+ if hasattr(self, "token_type_ids"):
+ buffered_token_type_ids = self.token_type_ids[:, :seq_length]
+ buffered_token_type_ids_expanded = buffered_token_type_ids.expand(input_shape[0], seq_length)
+ token_type_ids = buffered_token_type_ids_expanded
+ else:
+ token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+ if inputs_embeds is None:
+ inputs_embeds = self.word_embeddings(input_ids)
+ token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+ embeddings = inputs_embeds + token_type_embeddings
+ if self.position_embedding_type == "absolute":
+ position_embeddings = self.position_embeddings(position_ids)
+ embeddings += position_embeddings
+ # Add special node embedding
+ #embeddings[:, -2, :] = special_node_embedding
+ embeddings[torch.arange(embeddings.size(0)), spt] = special_node_embedding
+ embeddings = self.LayerNorm(embeddings)
+ embeddings = self.dropout(embeddings)
+ return embeddings
+
+
+class BertSelfAttention(nn.Module):
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+ raise ValueError(
+ f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
+ f"heads ({config.num_attention_heads})"
+ )
+
+ self.num_attention_heads = config.num_attention_heads
+ self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(config.hidden_size, self.all_head_size)
+ self.key = nn.Linear(config.hidden_size, self.all_head_size)
+ self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+ self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+ self.position_embedding_type = position_embedding_type or getattr(
+ config, "position_embedding_type", "absolute"
+ )
+ if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+ self.max_position_embeddings = config.max_position_embeddings
+ self.distance_embedding = nn.Embedding(2 * config.max_position_embeddings - 1, self.attention_head_size)
+
+ self.is_decoder = config.is_decoder
+
+ def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+ x = x.view(new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ head_mask: Optional[torch.FloatTensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
+ past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ output_attentions: Optional[bool] = False,
+ ) -> Tuple[torch.Tensor]:
+ mixed_query_layer = self.query(hidden_states)
+
+ # If this is instantiated as a cross-attention module, the keys
+ # and values come from an encoder; the attention mask needs to be
+ # such that the encoder's padding tokens are not attended to.
+ is_cross_attention = encoder_hidden_states is not None
+
+ if is_cross_attention and past_key_value is not None:
+ # reuse k,v, cross_attentions
+ key_layer = past_key_value[0]
+ value_layer = past_key_value[1]
+ attention_mask = encoder_attention_mask
+ elif is_cross_attention:
+ key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
+ value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
+ attention_mask = encoder_attention_mask
+ elif past_key_value is not None:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+ key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
+ value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
+ else:
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+
+ use_cache = past_key_value is not None
+ if self.is_decoder:
+ # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
+ # Further calls to cross_attention layer can then reuse all cross-attention
+ # key/value_states (first "if" case)
+ # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
+ # all previous decoder key/value_states. Further calls to uni-directional self-attention
+ # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
+ # if encoder bi-directional self-attention `past_key_value` is always `None`
+ past_key_value = (key_layer, value_layer)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ if self.position_embedding_type == "relative_key" or self.position_embedding_type == "relative_key_query":
+ query_length, key_length = query_layer.shape[2], key_layer.shape[2]
+ if use_cache:
+ position_ids_l = torch.tensor(key_length - 1, dtype=torch.long, device=hidden_states.device).view(
+ -1, 1
+ )
+ else:
+ position_ids_l = torch.arange(query_length, dtype=torch.long, device=hidden_states.device).view(-1, 1)
+ position_ids_r = torch.arange(key_length, dtype=torch.long, device=hidden_states.device).view(1, -1)
+ distance = position_ids_l - position_ids_r
+
+ positional_embedding = self.distance_embedding(distance + self.max_position_embeddings - 1)
+ positional_embedding = positional_embedding.to(dtype=query_layer.dtype) # fp16 compatibility
+
+ if self.position_embedding_type == "relative_key":
+ relative_position_scores = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+ attention_scores = attention_scores + relative_position_scores
+ elif self.position_embedding_type == "relative_key_query":
+ relative_position_scores_query = torch.einsum("bhld,lrd->bhlr", query_layer, positional_embedding)
+ relative_position_scores_key = torch.einsum("bhrd,lrd->bhlr", key_layer, positional_embedding)
+ attention_scores = attention_scores + relative_position_scores_query + relative_position_scores_key
+
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+ if attention_mask is not None:
+ # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+
+ # Mask heads if we want to
+ if head_mask is not None:
+ attention_probs = attention_probs * head_mask
+
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+
+ if self.is_decoder:
+ outputs = outputs + (past_key_value,)
+ return outputs
+
+
+class BertSelfOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertAttention(nn.Module):
+ def __init__(self, config, position_embedding_type=None):
+ super().__init__()
+ self.self = BertSelfAttention(config, position_embedding_type=position_embedding_type)
+ self.output = BertSelfOutput(config)
+ self.pruned_heads = set()
+
+ def prune_heads(self, heads):
+ if len(heads) == 0:
+ return
+ heads, index = find_pruneable_heads_and_indices(
+ heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+ )
+
+ # Prune linear layers
+ self.self.query = prune_linear_layer(self.self.query, index)
+ self.self.key = prune_linear_layer(self.self.key, index)
+ self.self.value = prune_linear_layer(self.self.value, index)
+ self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+ # Update hyper params and store pruned heads
+ self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+ self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+ self.pruned_heads = self.pruned_heads.union(heads)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ head_mask: Optional[torch.FloatTensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
+ past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ output_attentions: Optional[bool] = False,
+ ) -> Tuple[torch.Tensor]:
+ self_outputs = self.self(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ )
+ attention_output = self.output(self_outputs[0], hidden_states)
+ outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
+ return outputs
+
+
+class BertIntermediate(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+ if isinstance(config.hidden_act, str):
+ self.intermediate_act_fn = ACT2FN[config.hidden_act]
+ else:
+ self.intermediate_act_fn = config.hidden_act
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.intermediate_act_fn(hidden_states)
+ return hidden_states
+
+
+class BertOutput(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+ self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+ self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+ def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.dense(hidden_states)
+ hidden_states = self.dropout(hidden_states)
+ hidden_states = self.LayerNorm(hidden_states + input_tensor)
+ return hidden_states
+
+
+class BertLayer(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.chunk_size_feed_forward = config.chunk_size_feed_forward
+ self.seq_len_dim = 1
+ self.attention = BertAttention(config)
+ self.is_decoder = config.is_decoder
+ self.add_cross_attention = config.add_cross_attention
+ if self.add_cross_attention:
+ if not self.is_decoder:
+ raise ValueError(f"{self} should be used as a decoder model if cross attention is added")
+ self.crossattention = BertAttention(config, position_embedding_type="absolute")
+ self.intermediate = BertIntermediate(config)
+ self.output = BertOutput(config)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ head_mask: Optional[torch.FloatTensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
+ past_key_value: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ output_attentions: Optional[bool] = False,
+ ) -> Tuple[torch.Tensor]:
+ # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
+ self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
+ self_attention_outputs = self.attention(
+ hidden_states,
+ attention_mask,
+ head_mask,
+ output_attentions=output_attentions,
+ past_key_value=self_attn_past_key_value,
+ )
+ attention_output = self_attention_outputs[0]
+
+ # if decoder, the last output is tuple of self-attn cache
+ if self.is_decoder:
+ outputs = self_attention_outputs[1:-1]
+ present_key_value = self_attention_outputs[-1]
+ else:
+ outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
+
+ cross_attn_present_key_value = None
+ if self.is_decoder and encoder_hidden_states is not None:
+ if not hasattr(self, "crossattention"):
+ raise ValueError(
+ f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers"
+ " by setting `config.add_cross_attention=True`"
+ )
+
+ # cross_attn cached key/values tuple is at positions 3,4 of past_key_value tuple
+ cross_attn_past_key_value = past_key_value[-2:] if past_key_value is not None else None
+ cross_attention_outputs = self.crossattention(
+ attention_output,
+ attention_mask,
+ head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ cross_attn_past_key_value,
+ output_attentions,
+ )
+ attention_output = cross_attention_outputs[0]
+ outputs = outputs + cross_attention_outputs[1:-1] # add cross attentions if we output attention weights
+
+ # add cross-attn cache to positions 3,4 of present_key_value tuple
+ cross_attn_present_key_value = cross_attention_outputs[-1]
+ present_key_value = present_key_value + cross_attn_present_key_value
+
+ layer_output = apply_chunking_to_forward(
+ self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+ )
+ outputs = (layer_output,) + outputs
+
+ # if decoder, return the attn key/values as the last output
+ if self.is_decoder:
+ outputs = outputs + (present_key_value,)
+
+ return outputs
+
+ def feed_forward_chunk(self, attention_output):
+ intermediate_output = self.intermediate(attention_output)
+ layer_output = self.output(intermediate_output, attention_output)
+ return layer_output
+
+
+class BertEncoder(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.config = config
+ self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
+ self.gradient_checkpointing = False
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ head_mask: Optional[torch.FloatTensor] = None,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ encoder_attention_mask: Optional[torch.FloatTensor] = None,
+ past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
+ use_cache: Optional[bool] = None,
+ output_attentions: Optional[bool] = False,
+ output_hidden_states: Optional[bool] = False,
+ return_dict: Optional[bool] = True,
+ ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPastAndCrossAttentions]:
+ all_hidden_states = () if output_hidden_states else None
+ all_self_attentions = () if output_attentions else None
+ all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
+
+ if self.gradient_checkpointing and self.training:
+ if use_cache:
+ logger.warning_once(
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
+ )
+ use_cache = False
+
+ next_decoder_cache = () if use_cache else None
+ for i, layer_module in enumerate(self.layer):
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ layer_head_mask = head_mask[i] if head_mask is not None else None
+ past_key_value = past_key_values[i] if past_key_values is not None else None
+
+ if self.gradient_checkpointing and self.training:
+ layer_outputs = self._gradient_checkpointing_func(
+ layer_module.__call__,
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ )
+ else:
+ layer_outputs = layer_module(
+ hidden_states,
+ attention_mask,
+ layer_head_mask,
+ encoder_hidden_states,
+ encoder_attention_mask,
+ past_key_value,
+ output_attentions,
+ )
+
+ hidden_states = layer_outputs[0]
+ if use_cache:
+ next_decoder_cache += (layer_outputs[-1],)
+ if output_attentions:
+ all_self_attentions = all_self_attentions + (layer_outputs[1],)
+ if self.config.add_cross_attention:
+ all_cross_attentions = all_cross_attentions + (layer_outputs[2],)
+
+ if output_hidden_states:
+ all_hidden_states = all_hidden_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [
+ hidden_states,
+ next_decoder_cache,
+ all_hidden_states,
+ all_self_attentions,
+ all_cross_attentions,
+ ]
+ if v is not None
+ )
+ return BaseModelOutputWithPastAndCrossAttentions(
+ last_hidden_state=hidden_states,
+ past_key_values=next_decoder_cache,
+ hidden_states=all_hidden_states,
+ attentions=all_self_attentions,
+ cross_attentions=all_cross_attentions,
+ )
+
+
+class BertPooler(nn.Module):
+ def __init__(self, config):
+ super().__init__()
+ self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+ self.activation = nn.Tanh()
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ # We "pool" the model by simply taking the hidden state corresponding
+ # to the first token.
+ first_token_tensor = hidden_states[:, 0]
+ pooled_output = self.dense(first_token_tensor)
+ pooled_output = self.activation(pooled_output)
+ return pooled_output
+
+
+
+class GreaseLM(torch.nn.Module):
+ def __init__(self, args, input_dim=768, hidden_channels=128, layers=5):
+ super(GreaseLM, self).__init__()
+ self.args = args
+ self.num_relations = 12
+ text_layers = 4
+
+ def forward(self, x, edge_index, edge_type, batch, inds, special_node_idx=None):
+ for layer in self.convs[:-1]:
+ x = F.relu(layer(x, edge_index, edge_type.view(-1)))
+ x = self.convs[-1](x, edge_index, edge_type.view(-1))
+ if self.args.fusion:
+ special_mask = special_node_idx.view(-1, 1)
+ z = x[special_mask[:, 0]]
+ inds = torch.stack(inds, dim=-1).tolist()
+ for i in range(1, len(inds)):
+ v_c = 0
+ for j in range(i, 0, -1):
+ v_c += batch.tolist().count(j-1)
+ inds[i][0] += v_c
+ inds[i][1] += v_c
+ mask = torch.tensor(inds)
+ head = x[mask[:, 0]]
+ tail = x[mask[:, 1]]
+ avg_n = torch.cat([z, head, tail], dim=-1).view(special_mask.shape[0], 3, -1)
+ pool = torch.mean(avg_n, dim=1)
+ #pool = global_mean_pool(x, batch)
+ return pool, z
+ else:
+ inds = torch.stack(inds, dim=-1).tolist()
+ for i in range(1, len(inds)):
+ v_c = 0
+ for j in range(i, 0, -1):
+ v_c += batch.tolist().count(j-1)
+ inds[i][0] += v_c
+ inds[i][1] += v_c
+ mask = torch.tensor(inds)
+ head = x[mask[:, 0]]
+ tail = x[mask[:, 1]]
+ x = torch.cat([head, tail], dim=-1)
+ return x
diff --git a/dataset.py b/dataset.py
@@ -0,0 +1,148 @@
+from device import device
+from torch.utils.data import Dataset
+from torch_geometric.data import Data
+from torch_geometric.utils.convert import from_scipy_sparse_matrix
+from tqdm import trange
+from transformers import AutoTokenizer, AutoModel
+from typing import List, Dict
+from utils import load_graph_from_path
+import json
+import networkx as nx
+import torch
+PLM = 'bert-base-uncased'
+
+
+class TextGraphDataset(Dataset):
+ def __init__(self, data, args, head_emb=None, tail_emb=None, tokenizer=None):
+ self.special_token = "[SPECIAL_TOKEN]"
+ self.special_node = "SPECIAL_NODE]"
+ self.questions: List[str] = []
+ self.question_paths: List[List[str]] = []
+ self.paths: List[List[List[str], List[str]]] = []
+ self.labels: List[int] = []
+ self.sample_ids: List[str] = []
+ self.graphs: List[Data] = []
+ self.indicies: List[List[int]] = []
+ self.unidirectional: bool = args.unidirectional
+ self.bidirectional: bool = args.bidirectional
+ self.model = AutoModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device)
+ node_embedding_size: int = self.model.config.hidden_size
+
+ if tokenizer is None:
+ self.tokenizer = AutoTokenizer.from_pretrained(PLM, use_fast=False)
+ if args.bidirectional:
+ assert args.unidirectional, "Unidirectional is set to false!"
+ self.tokenizer.add_tokens(self.special_token, special_tokens=True)
+ self.special_token_id = len(self.tokenizer) - 1
+ else:
+ self.tokenizer = tokenizer
+
+ self.texts: Dict[str, List[int]] = {
+ "input_ids": [],
+ "attention_mask": []
+ }
+
+ if self.bidirectional:
+ self.model.resize_token_embeddings(len(self.tokenizer))
+
+ if self.unidirectional:
+ self.special_node_indicies: List[List[int, int]] = []
+
+ if head_emb is None and tail_emb is None:
+ self.head_emb: torch.Tensor = torch.randn(node_embedding_size)
+ self.tail_emb: torch.Tensor = torch.randn(node_embedding_size)
+ else:
+ self.head_emb: torch.Tensor = head_emb
+ self.tail_emb: torch.Tensor = tail_emb
+
+ with open("static/edge_type_dict.json", 'r') as f:
+ e2i = json.load(f)
+
+ for i in trange(len(data)):
+ graph_path: str = data[i]["graph"]
+ raw_graph: nx.MultiDiGraph = load_graph_from_path(graph_path)
+ question: str = data[i]["question"]
+ self.labels.append(data[i]["label"])
+ self.sample_ids.append(i)
+ self.questions.append(question)
+ self.question_paths.append(data[i]["question_path"])
+ self.paths.append(data[i]["og_path"])
+
+ if self.bidirectional:
+ question += self.special_token
+
+ tokenized_question = self.tokenizer(question, padding="max_length", max_length=64)
+ question_input_ids = torch.tensor(tokenized_question["input_ids"]).view(1, -1).to(device)
+ question_attention_mask = torch.tensor(tokenized_question["attention_mask"]).view(1, -1).to(device)
+ self.texts["input_ids"].append(question_input_ids[0])
+ self.texts["attention_mask"].append(question_attention_mask[0])
+
+ head_string = data[i]["og_path"][0][0]
+ tail_string = data[i]["og_path"][-1][-1]
+ if args.GNN_only or args.static:
+ head = self.head_emb.tolist()
+ tail = self.tail_emb.tolist()
+ else:
+ head_tokenized = self.tokenizer(question, head_string, padding="max_length", max_length=64)
+ tail_tokenized = self.tokenizer(question, tail_string, padding="max_length", max_length=64)
+ head_input_ids = torch.tensor(head_tokenized["input_ids"]).view(1, -1).to(device)
+ tail_input_ids = torch.tensor(tail_tokenized["input_ids"]).view(1, -1).to(device)
+ head_attention_mask = torch.tensor(head_tokenized["attention_mask"]).view(1, -1).to(device)
+ tail_attention_mask = torch.tensor(tail_tokenized["attention_mask"]).view(1, -1).to(device)
+ with torch.no_grad():
+ head_output = self.model(head_input_ids, head_attention_mask)[1].squeeze(0)
+ tail_output = self.model(tail_input_ids, tail_attention_mask)[1].squeeze(0)
+ head = head_output.tolist()
+ tail = tail_output.tolist()
+
+ if self.unidirectional or self.bidirectional:
+ with torch.no_grad():
+ question_embedded = self.model(question_input_ids, question_attention_mask)[1].squeeze(0)
+ raw_graph.add_node(self.special_node)
+ for node in raw_graph.nodes():
+ if node == self.special_node:
+ continue
+ raw_graph.add_edge(self.special_node, node, type="special", id="special")
+ raw_graph.add_edge(node, self.special_node, type="special_r", id="special_r")
+
+ edge_type_tensor = torch.tensor([e2i[e[-1]["type"]] for e in raw_graph.edges(data=True)], dtype=torch.long)
+ adj_matrix = nx.adjacency_matrix(raw_graph)
+ edge_index = from_scipy_sparse_matrix(adj_matrix)
+ node_embeddings = {n: torch.zeros(node_embedding_size).tolist() for i, n in enumerate(list(raw_graph.nodes()))}
+ node_embeddings[head_string] = head
+ node_embeddings[tail_string] = tail
+ node_keys = list(node_embeddings.keys())
+ head_index = node_keys.index(head_string)
+ tail_index = node_keys.index(tail_string)
+ if self.unidirectional or self.bidirectional:
+ node_embeddings[self.special_node] = question_embedded
+ special_node_index = node_keys.index(self.special_node)
+ self.special_node_indicies.append(special_node_index)
+ nodes = torch.tensor(list(node_embeddings.values()), dtype=torch.float)
+ processed_graph = Data(x=nodes, edge_index=edge_index[0], edge_type=edge_type_tensor).to(device)
+ self.graphs.append(processed_graph)
+ self.indicies.append([head_index, tail_index])
+
+ def __len__(self):
+ return len(self.labels)
+
+ def __getitem__(self, idx):
+ input_ids = self.texts["input_ids"][idx]
+ attention_mask = self.texts["attention_mask"][idx]
+ label = self.labels[idx]
+
+ if self.unidirectional or self.bidirectional:
+ special_node_index = self.special_node_indicies
+ else:
+ special_node_index = [0 for _ in self.indicies]
+
+ if self.bidirectional:
+ token = self.tokenizer(self.special_token, add_special_tokens=False)["input_ids"][0]
+ special_token_index = input_ids.tolist().index(token)
+ else:
+ special_token_index = [0 for _ in self.indicies]
+
+ 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)
+
+ def decode(self, idx):
+ return self.tokenizer.decode(self.texts["input_ids"][idx], skip_special_tokens=False), self.labels[idx]
diff --git a/device.py b/device.py
@@ -0,0 +1,2 @@
+import torch
+device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
diff --git a/model.py b/model.py
@@ -0,0 +1,116 @@
+from device import device
+from bidirectional import BertModel
+from torch_geometric.nn import RGCNConv
+from transformers import AutoModel
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+PLM = 'bert-base-uncased'
+
+
+class RelationalGCN(torch.nn.Module):
+ def __init__(self, args, input_dim=768, hidden_channels=128,
+ layers=5):
+ super(RelationalGCN, self).__init__()
+ self.args = args
+ num_relations = 12 if self.args.unidirectional or self.args.bidirectional else 10
+ self.convs = torch.nn.ModuleList(
+ [RGCNConv(in_channels=input_dim, out_channels=hidden_channels, num_relations=num_relations)] +
+ [RGCNConv(in_channels=hidden_channels, out_channels=hidden_channels, num_relations=num_relations)
+ for _ in range(layers-2)] +
+ [RGCNConv(in_channels=hidden_channels, out_channels=hidden_channels, num_relations=num_relations)]
+ )
+
+ def forward(self, x, edge_index, edge_type, batch, inds, special_node_idx=None):
+ for layer in self.convs[:-1]:
+ x = F.relu(layer(x, edge_index, edge_type.view(-1)))
+ x = self.convs[-1](x, edge_index, edge_type.view(-1))
+ if self.args.unidirectional:
+ special_mask = special_node_idx.view(-1, 1)
+ z = x[special_mask[:, 0]]
+ inds = torch.stack(inds, dim=-1).tolist()
+ for i in range(1, len(inds)):
+ v_c = 0
+ for j in range(i, 0, -1):
+ v_c += batch.tolist().count(j-1)
+ inds[i][0] += v_c
+ inds[i][1] += v_c
+ mask = torch.tensor(inds)
+ head = x[mask[:, 0]]
+ tail = x[mask[:, 1]]
+ avg_n = torch.cat([z, head, tail], dim=-1).view(special_mask.shape[0], 3, -1)
+ pool = torch.mean(avg_n, dim=1)
+ return pool, z
+ else:
+ inds = torch.stack(inds, dim=-1).tolist()
+ for i in range(1, len(inds)):
+ v_c = 0
+ for j in range(i, 0, -1):
+ v_c += batch.tolist().count(j-1)
+ inds[i][0] += v_c
+ inds[i][1] += v_c
+ mask = torch.tensor(inds)
+ head = x[mask[:, 0]]
+ tail = x[mask[:, 1]]
+ x = torch.cat([head, tail], dim=-1)
+ return x
+
+
+class TextRGCN(torch.nn.Module):
+ def __init__(self, args, graph_input_dim=768):
+ super(TextRGCN, self).__init__()
+ self.hidden_channels = args.gnn_hidden
+ self.GNN_only = args.GNN_only
+ self.LM_only = args.LM_only
+ self.gcn = RelationalGCN(args, graph_input_dim, self.hidden_channels, layers=args.gnn_layers).to(device)
+ self.dropout_t = nn.Dropout(args.dropout)
+ self.dropout_g = nn.Dropout(args.dropout)
+ self.unidirectional = args.unidirectional
+ self.bidirectional = args.bidirectional
+ if args.bidirectional:
+ self.text_encoder = BertModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device)
+ else:
+ self.text_encoder = AutoModel.from_pretrained(PLM, return_dict=True, output_hidden_states=True).to(device)
+
+ gnn_output_dim = self.hidden_channels if self.unidirectional else self.hidden_channels * 2
+
+ if self.GNN_only:
+ self.output = nn.Linear(gnn_output_dim, gnn_output_dim)
+ self.output2 = nn.Linear(gnn_output_dim, 1)
+ elif self.LM_only:
+ self.output = nn.Linear(self.text_encoder.config.hidden_size, self.text_encoder.config.hidden_size,)
+ self.output2 = nn.Linear(self.text_encoder.config.hidden_size, 1)
+ elif self.unidirectional:
+ self.output = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim * 2, self.text_encoder.config.hidden_size,)
+ self.output2 = nn.Linear(self.text_encoder.config.hidden_size, 1)
+ else:
+ self.output = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim, self.text_encoder.config.hidden_size + gnn_output_dim)
+ self.output2 = nn.Linear(self.text_encoder.config.hidden_size + gnn_output_dim, 1)
+
+ def forward(self, input_ids, attention_maskask, data, inds, special_node_indicies, special_token_indicies):
+ batch = data.batch.to(device)
+ if self.GNN_only:
+ enc = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds)
+ enc = self.dropout_g(enc)
+ elif self.LM_only:
+ enc = self.text_encoder(input_ids, attention_maskask)[1]
+ enc = self.dropout_t(enc)
+ else:
+ if self.bidirectional:
+ pooled, special_node_embedding = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=special_node_indicies)
+ z_text = self.text_encoder(input_ids, attention_maskask, special_node_embedding=special_node_embedding, spt=special_token_indicies)[1]
+ z_text = self.dropout_t(z_text)
+ enc = torch.cat([z_text, special_node_embedding, pooled], dim=-1)
+ elif self.unidirectional:
+ pooled, special_node_embedding = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=special_node_indicies)
+ z_text = self.text_encoder(input_ids, attention_maskask)[1]
+ z_text = self.dropout_t(z_text)
+ enc = torch.cat([z_text, special_node_embedding, pooled], dim=-1)
+ else: # The else statement is entered for the Disjoint and Grounded models
+ z_graph = self.gcn(data.x, data.edge_index, data.edge_type, batch, inds, special_node_idx=None)
+ z_text = self.text_encoder(input_ids, attention_maskask)[1]
+ z_text = self.dropout_t(z_text)
+ enc = torch.cat([z_graph, z_text], dim=-1)
+ out = F.relu(self.output(enc))
+ out = self.output2(out)
+ return out
diff --git a/train.py b/train.py
@@ -0,0 +1,221 @@
+from dataset import TextGraphDataset
+from device import device
+from model import TextRGCN
+from torch_geometric.loader import DataLoader
+from tqdm import tqdm
+import argparse
+import jsonlines
+import logging
+import math
+import numpy as np
+import random
+import torch
+import torch.nn as nn
+from utils import set_seed
+import wandb
+
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Compositional reasoning over KGs and text using LM+GNNs")
+ parser.add_argument("--batch_size", type=int, default=16, help="The batch size to use during training.")
+ parser.add_argument("--weight_decay", type=float, default=0.001, help="The weight decay to use during training.")
+ parser.add_argument("--model_name", type=str, default="bert-base-uncased", help="The pretrained model to use")
+ parser.add_argument("--scheduler", type=str, default="constant", help="The lr scheduler protocol to use")
+ parser.add_argument("--train_path", type=str, default=None, help="The question dataset file for train")
+ parser.add_argument("--val_path", type=str, default=None, help="The question dataset file for val")
+ parser.add_argument("--output_path", type=str, default=None, help="The output predictions path")
+ parser.add_argument("--epochs", type=int, default=3, help="The number of epochs.")
+ parser.add_argument("--gnn_layers", type=int, default=5, help="The number of layers in the GNN model.")
+ parser.add_argument("--gnn_hidden", type=int, default=1024, help="The hidden size of the GNN model.")
+ parser.add_argument("--static", action='store_true', help="Use static entity embeddings")
+ parser.add_argument("--unidirectional", action='store_true', help="Use deep unidirectional mode")
+ parser.add_argument("--bidirectional", action='store_true', help="Use deep unidirectional mode")
+ parser.add_argument("--debug", action='store_true', help="Trigger debug mode")
+ parser.add_argument("--GNN_only", action='store_true', help="Trigger GNN only ablation")
+ parser.add_argument("--LM_only", action='store_true', help="Trigger LM only ablation")
+ parser.add_argument("--lr", type=float, default=3e-5, help="The learning rate).")
+ parser.add_argument("--gnn_lr", type=float, default=3e-5, help="The learning rate).")
+ parser.add_argument("--llm_lr", type=float, default=1e-6, help="The learning rate).")
+ parser.add_argument("--head_lr", type=float, default=1e-5, help="The learning rate).")
+ #parser.add_argument("--seed", type=int, default=42, help="The rng seed")
+ parser.add_argument('--seeds', nargs='+', help='Set of seeds to run with', required=True)
+ parser.add_argument("--device_max_steps", default=200000, type=int, help="Total number of training steps to perform per device")
+ parser.add_argument("--warmup_proportion", default=0.016, type=float, help="Proportion of training to perform linear learning rate warmup for. E.g., 0.1 = 10%% of training.")
+ parser.add_argument("--dropout", type=float, default=0.2, help="The dropout value")
+ args = parser.parse_args()
+ return args
+
+
+def train(model, loader):
+ total_loss = 0.0
+ correct = 0.0
+ n = 0.0
+ for uids, input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index, label in tqdm(loader, leave=False):
+ optimizer.zero_grad()
+ input_ids = input_ids.to(device)
+ attention_mask = attention_mask.to(device)
+ label = label.to(device)
+ out = model(input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index)
+ loss = criterion(out.squeeze(-1), label)
+ total_loss += loss.item()
+ pred = (nn.functional.sigmoid(out) > 0.5).float().squeeze()
+ correct += (pred == label).sum()
+ n += len(label)
+ loss.backward()
+ optimizer.step()
+ scheduler.step()
+ acc = correct / n
+ loss = total_loss / len(loader)
+ return acc, loss
+
+
+def val(model, loader):
+ val_loss = 0.0
+ val_preds = []
+ val_raw_preds = []
+ val_uids = []
+ correct = 0.0
+ n = 0.0
+ with torch.no_grad():
+ for uids, input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index, label in tqdm(val_loader, leave=False):
+ input_ids = input_ids.to(device)
+ attention_mask = attention_mask.to(device)
+ label = label.to(device)
+ out = model(input_ids, attention_mask, graph_data, head_tail_indicies, special_node_index, special_token_index)
+ loss = criterion(out.squeeze(-1), label)
+ val_loss += loss.item()
+ pred = (nn.functional.sigmoid(out) > 0.5).float().squeeze()
+ correct += (pred == label).sum()
+ n += len(label)
+ val_uids += uids.tolist()
+ val_raw_preds += out.cpu().tolist()
+ val_preds += pred.tolist()
+ val_acc = correct / n
+ val_loss = val_loss / len(val_loader)
+ return val_acc, val_loss, val_uids, val_raw_preds, val_preds
+
+
+if __name__ == '__main__':
+ args = parse_args()
+ logging.basicConfig(
+ format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+ datefmt="%m/%d/%Y %H:%M:%S",
+ level=logging.INFO,
+ )
+ set_seed(0)
+ with jsonlines.open(args.train_path, 'r') as f:
+ train_questions = [o for o in f]
+ with jsonlines.open(args.val_path, 'r') as f:
+ val_questions = [o for o in f]
+
+ train_dataset = TextGraphDataset(train_questions, args=args)
+ val_dataset = TextGraphDataset(val_questions,
+ args=args,
+ head_emb=train_dataset.head_emb,
+ tail_emb=train_dataset.tail_emb,
+ tokenizer=train_dataset.tokenizer)
+
+ logging.info(f'{train_dataset.decode(0)}')
+ logging.info('Starting training...')
+ for seed in args.seeds:
+ logging.info(f'Setting seed {seed}...')
+ set_seed(int(seed))
+ train_loader = DataLoader(train_dataset,
+ batch_size=args.batch_size,
+ shuffle=True,
+ drop_last=True)
+ val_loader = DataLoader(val_dataset,
+ batch_size=args.batch_size,
+ shuffle=False,
+ drop_last=False)
+ config = {
+ 'lr': args.lr,
+ 'gnn_lr': args.gnn_lr,
+ 'lm_lr': args.llm_lr,
+ 'head_lr': args.head_lr,
+ 'epochs': args.epochs,
+ 'batch_size': args.batch_size,
+ 'weight_decay': args.weight_decay,
+ 'dropout': args.dropout,
+ 'gnn_layers': args.gnn_layers,
+ 'scheduler': args.scheduler,
+ 'LM_ONLY': args.LM_only,
+ 'GNN_ONLY': args.GNN_only,
+ 'unidirectional': args.unidirectional,
+ 'bidirectional': args.bidirectional,
+ 'disjoint': args.static,
+ }
+ wandb.init(project="results", config=config, entity="sondrewo", name=f"{args.output_path}_seed_{seed}")
+ criterion = torch.nn.BCEWithLogitsLoss()
+ model = TextRGCN(args).to(device)
+ if args.bidirectional:
+ model.text_encoder.resize_token_embeddings(len(train_dataset.tokenizer))
+
+ optimizer = torch.optim.AdamW(
+ [
+ {"params": model.gcn.parameters(), "lr": args.gnn_lr},
+ {"params": model.text_encoder.parameters(), "lr": args.llm_lr},
+ {"params": model.output.parameters(), "lr": args.head_lr},
+ {"params": model.output2.parameters(), "lr": args.head_lr},
+ ],
+ lr=args.lr,
+ betas=(0.9, 0.98),
+ eps=1e-6,
+ )
+
+ args.device_max_steps = args.epochs * len(train_loader)
+ if args.scheduler == "cosine":
+ def cosine_schedule_with_warmup(optimizer, num_warmup_steps: int, num_training_steps: int, min_factor: float):
+ def lr_lambda(current_step):
+ if current_step < num_warmup_steps:
+ return float(current_step) / float(max(1, num_warmup_steps))
+ progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))
+ lr = max(min_factor, min_factor + (1 - min_factor) * 0.5 * (1.0 + math.cos(math.pi * progress)))
+ return lr
+ return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+ scheduler = cosine_schedule_with_warmup(optimizer, int(args.device_max_steps * args.warmup_proportion), args.device_max_steps, 0.1)
+
+ elif args.scheduler == "linear":
+ scheduler = torch.optim.lr_scheduler.ChainedScheduler([
+ torch.optim.lr_scheduler.LinearLR(optimizer,
+ start_factor=1e-9,
+ end_factor=1.0,
+ total_iters=int(args.device_max_steps * args.warmup_proportion)),
+ torch.optim.lr_scheduler.LinearLR(optimizer,
+ start_factor=1.0,
+ end_factor=1e-9,
+ total_iters=args.device_max_steps)
+ ])
+
+ elif args.scheduler == "constant":
+ scheduler = torch.optim.lr_scheduler.ConstantLR(optimizer, factor=1.0, total_iters=1)
+
+ for epoch in range(args.epochs):
+ model.train()
+ train_acc, train_loss = train(model, train_loader)
+ wandb.log({'training_accuracy': train_acc, 'training_loss': train_loss})
+ logging.info(f'Train accuracy: {train_acc}, train.loss: {train_loss}')
+
+ model.eval()
+ val_acc, val_loss, val_uids, val_raw_preds, val_preds = val(model, val_loader)
+ wandb.log({'val_accuracy': val_acc, 'val_loss': val_loss})
+ logging.info(f'val accuracy: {val_acc}, val.loss: {val_loss}')
+
+ logging.info('Writing final predictions...')
+ final_preds = []
+ for uid, raw, pred in zip(val_uids, val_raw_preds, val_preds):
+ obj = {
+ 'uid': uid,
+ 'question': val_dataset.questions[uid],
+ 'relations': val_dataset.question_paths[uid],
+ 'head_tail': val_dataset.paths[uid],
+ 'label': val_dataset.labels[uid],
+ 'raw_pred': raw,
+ 'pred': pred
+ }
+ final_preds.append(obj)
+
+ with jsonlines.open(f'./results/{args.output_path}_seed_{seed}.jsonl', 'w') as f:
+ f.write_all(final_preds)
+ print('Saved predictions...')
+ wandb.finish()
diff --git a/utils.py b/utils.py
@@ -0,0 +1,27 @@
+import networkx as nx
+import torch
+import numpy as np
+import random
+
+def load_graph_from_path(path: str) -> nx.Graph:
+ '''
+ Load a graph from a given path.
+
+ Args:
+ path (str): The path to the graph file.
+
+ Returns:
+ G (networkx.Graph): The loaded graph.
+
+ '''
+ G = nx.read_graphml(f'{path}')
+ return G
+
+def set_seed(seed: int) -> None:
+ torch.manual_seed(seed)
+ torch.cuda.manual_seed(seed)
+ torch.cuda.manual_seed_all(seed)
+ np.random.seed(seed)
+ random.seed(seed)
+ torch.backends.cudnn.benchmark = False
+ torch.backends.cudnn.deterministic = True
+\ No newline at end of file
