From d73c52e15b519af764a83378d4eab19fb31985e0 Mon Sep 17 00:00:00 2001 From: Gustaf Rydholm Date: Fri, 30 Sep 2022 00:33:53 +0200 Subject: Update transformer --- text_recognizer/networks/transformer/decoder_block.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) (limited to 'text_recognizer/networks/transformer/decoder_block.py') diff --git a/text_recognizer/networks/transformer/decoder_block.py b/text_recognizer/networks/transformer/decoder_block.py index b80a078..2dc4ddf 100644 --- a/text_recognizer/networks/transformer/decoder_block.py +++ b/text_recognizer/networks/transformer/decoder_block.py @@ -5,6 +5,7 @@ from typing import Optional, Type from torch import Tensor, nn from text_recognizer.networks.transformer.attention import Attention +from text_recognizer.networks.transformer.embeddings.rotary import RotaryEmbedding from text_recognizer.networks.transformer.ff import FeedForward @@ -25,22 +26,19 @@ class DecoderBlock(nn.Module): self.cross_attn = cross_attn self.ln_ff = deepcopy(norm) self.ff = ff - self.has_pos_emb = self.attn.rotary_embedding is not None def forward( self, x: Tensor, context: Optional[Tensor] = None, - input_mask: Optional[Tensor] = None, - context_mask: Optional[Tensor] = None, + mask: Optional[Tensor] = None, + rotary_embedding: Optional[RotaryEmbedding] = None, ) -> Tensor: """Applies decoder block on input signals.""" - x = x + self.attn(self.ln_attn(x), input_mask=input_mask) + x = x + self.attn(self.ln_attn(x), mask=mask, rotary_embedding=rotary_embedding) x = x + self.cross_attn( x=self.ln_cross_attn(x), context=context, - input_mask=input_mask, - context_mask=context_mask, ) x = x + self.ff(self.ln_ff(x)) return x -- cgit v1.2.3-70-g09d2