rat-sql-gap/seq2struct/models/transformer.py [213:329]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - class MultiHeadedAttentionWithRelations(nn.Module): def __init__(self, h, d_model, dropout=0.1): "Take in model size and number of heads." super(MultiHeadedAttentionWithRelations, self).__init__() assert d_model % h == 0 # We assume d_v always equals d_k self.d_k = d_model // h self.h = h self.linears = clones(lambda: nn.Linear(d_model, d_model), 4) self.attn = None self.dropout = nn.Dropout(p=dropout) def forward(self, query, key, value, relation_k, relation_v, mask=None): # query shape: [batch, num queries, d_model] # key shape: [batch, num kv, d_model] # value shape: [batch, num kv, d_model] # relations_k shape: [batch, num queries, num kv, (d_model // h)] # relations_v shape: [batch, num queries, num kv, (d_model // h)] # mask shape: [batch, num queries, num kv] if mask is not None: # Same mask applied to all h heads. # mask shape: [batch, 1, num queries, num kv] mask = mask.unsqueeze(1) nbatches = query.size(0) # 1) Do all the linear projections in batch from d_model => h x d_k query, key, value = \ [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))] # 2) Apply attention on all the projected vectors in batch. # x shape: [batch, heads, num queries, depth] x, self.attn = attention_with_relations( query, key, value, relation_k, relation_v, mask=mask, dropout=self.dropout) # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous() \ .view(nbatches, -1, self.h * self.d_k) return self.linears[-1](x) # Adapted from The Annotated Transformer class Encoder(nn.Module): "Core encoder is a stack of N layers" def __init__(self, layer, layer_size, N, tie_layers=False): super(Encoder, self).__init__() if tie_layers: self.layer = layer() self.layers = [self.layer for _ in range(N)] else: self.layers = clones(layer, N) self.norm = nn.LayerNorm(layer_size) # TODO initialize using xavier def forward(self, x, relation, mask): "Pass the input (and mask) through each layer in turn." for layer in self.layers: x = layer(x, relation, mask) return self.norm(x) # Adapted from The Annotated Transformer class SublayerConnection(nn.Module): """ A residual connection followed by a layer norm. Note for code simplicity the norm is first as opposed to last. """ def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm = nn.LayerNorm(size) self.dropout = nn.Dropout(dropout) def forward(self, x, sublayer): "Apply residual connection to any sublayer with the same size." return x + self.dropout(sublayer(self.norm(x))) # Adapted from The Annotated Transformer class EncoderLayer(nn.Module): "Encoder is made up of self-attn and feed forward (defined below)" def __init__(self, size, self_attn, feed_forward, num_relation_kinds, dropout): super(EncoderLayer, self).__init__() self.self_attn = self_attn self.feed_forward = feed_forward self.sublayer = clones(lambda: SublayerConnection(size, dropout), 2) self.size = size self.relation_k_emb = nn.Embedding(num_relation_kinds, self.self_attn.d_k) self.relation_v_emb = nn.Embedding(num_relation_kinds, self.self_attn.d_k) def forward(self, x, relation, mask): "Follow Figure 1 (left) for connections." relation_k = self.relation_k_emb(relation) relation_v = self.relation_v_emb(relation) x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, relation_k, relation_v, mask)) return self.sublayer[1](x, self.feed_forward) # Adapted from The Annotated Transformer class PositionwiseFeedForward(nn.Module): "Implements FFN equation." def __init__(self, d_model, d_ff, dropout=0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 = nn.Linear(d_model, d_ff) self.w_2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.w_2(self.dropout(F.relu(self.w_1(x)))) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - relogic/pretrainkit/models/relationalsemparse/relational_transformer.py [225:345]: - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - class MultiHeadedAttentionWithRelations(nn.Module): def __init__(self, h, d_model, dropout=0.1): "Take in model size and number of heads." super(MultiHeadedAttentionWithRelations, self).__init__() assert d_model % h == 0 # We assume d_v always equals d_k self.d_k = d_model // h self.h = h self.linears = clones(lambda: nn.Linear(d_model, d_model), 4) self.attn = None self.dropout = nn.Dropout(p=dropout) def forward(self, query, key, value, relation_k, relation_v, mask=None): # query shape: [batch, num queries, d_model] # key shape: [batch, num kv, d_model] # value shape: [batch, num kv, d_model] # relations_k shape: [batch, num queries, num kv, (d_model // h)] # relations_v shape: [batch, num queries, num kv, (d_model // h)] # mask shape: [batch, num queries, num kv] if mask is not None: # Same mask applied to all h heads. # mask shape: [batch, 1, num queries, num kv] mask = mask.unsqueeze(1) nbatches = query.size(0) # 1) Do all the linear projections in batch from d_model => h x d_k query, key, value = \ [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2) for l, x in zip(self.linears, (query, key, value))] # 2) Apply attention on all the projected vectors in batch. # x shape: [batch, heads, num queries, depth] x, self.attn = attention_with_relations( query, key, value, relation_k, relation_v, mask=mask, dropout=self.dropout) # 3) "Concat" using a view and apply a final linear. x = x.transpose(1, 2).contiguous() \ .view(nbatches, -1, self.h * self.d_k) return self.linears[-1](x) # Adapted from The Annotated Transformer class Encoder(nn.Module): "Core encoder is a stack of N layers" def __init__(self, layer, layer_size, N, tie_layers=False): super(Encoder, self).__init__() if tie_layers: self.layer = layer() self.layers = [self.layer for _ in range(N)] else: self.layers = clones(layer, N) self.norm = nn.LayerNorm(layer_size) # TODO initialize using xavier def forward(self, x, relation, mask): "Pass the input (and mask) through each layer in turn." for layer in self.layers: x = layer(x, relation, mask) return self.norm(x) # Adapted from The Annotated Transformer class SublayerConnection(nn.Module): """ A residual connection followed by a layer norm. Note for code simplicity the norm is first as opposed to last. """ def __init__(self, size, dropout): super(SublayerConnection, self).__init__() self.norm = nn.LayerNorm(size) self.dropout = nn.Dropout(dropout) def forward(self, x, sublayer): "Apply residual connection to any sublayer with the same size." return x + self.dropout(sublayer(self.norm(x))) # Adapted from The Annotated Transformer class EncoderLayer(nn.Module): "Encoder is made up of self-attn and feed forward (defined below)" def __init__(self, size, self_attn, feed_forward, num_relation_kinds, dropout): super(EncoderLayer, self).__init__() self.self_attn = self_attn self.feed_forward = feed_forward self.sublayer = clones(lambda: SublayerConnection(size, dropout), 2) self.size = size self.relation_k_emb = nn.Embedding(num_relation_kinds, self.self_attn.d_k) self.relation_v_emb = nn.Embedding(num_relation_kinds, self.self_attn.d_k) def forward(self, x, relation, mask): "Follow Figure 1 (left) for connections." relation_k = self.relation_k_emb(relation) relation_v = self.relation_v_emb(relation) x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, relation_k, relation_v, mask)) return self.sublayer[1](x, self.feed_forward) # Adapted from The Annotated Transformer class PositionwiseFeedForward(nn.Module): "Implements FFN equation." def __init__(self, d_model, d_ff, dropout=0.1): super(PositionwiseFeedForward, self).__init__() self.w_1 = nn.Linear(d_model, d_ff) self.w_2 = nn.Linear(d_ff, d_model) self.dropout = nn.Dropout(dropout) def forward(self, x): return self.w_2(self.dropout(F.relu(self.w_1(x)))) - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -