<jupyter_start><jupyter_text>Models QRNN2QRNN><jupyter_code>enc_vocab_size = 21128
enc_pad_id = 0

dec_vocab_size = 50259
dec_pad_id = 50258

embeded_size = 512
num_encoder_layers = 2
num_decoder_layers = 2
drop_p = 0.1<jupyter_output><empty_output><jupyter_text>QRNNEncoder<jupyter_code># export
class QRNNEncoder(nn.Module):
    def __init__(
        self,
        vocab_size, embeded_size, pad_id,
        num_layers=1, drop_p=0,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embeded_size, padding_idx=pad_id)
        self.encoder = QRNN(embeded_size, embeded_size, n_layers=num_layers, dropout=drop_p, bidirectional=True, save_prev_x=False)
        
    def forward(self, src_inp_ids):
        '''
            src_inp_ids: (bs, enc_seq_len)
            returns: output, h
                output: (bs, seq_len, 2*embeded_size)
                h: (2*num_layers, bs, embeded_size)
        '''
        embeded = self.embedding(src_inp_ids) # (bs, enc_seq_len, embeded_size)
        output, h = self.encoder(embeded) # (bs, enc_seq_len, 2*embeded_size), (2*num_encoder_layers, bs, embeded_size)
        return output, h
src_input_ids = torch.randint(0, enc_vocab_size, (bs, enc_seq_len)) # (bs, enc_seq_len)
encoder = QRNNEncoder(enc_vocab_size, embeded_size, enc_pad_id, num_encoder_layers, drop_p)
output, h = encoder(src_input_ids)
test_eq(output.shape, (bs, enc_seq_len, 2*embeded_size))
test_eq(h.shape, (2*num_encoder_layers, bs, embeded_size))<jupyter_output><empty_output><jupyter_text>QRNNDecoder<jupyter_code># export
class QRNNDecoder(nn.Module):
    def __init__(
        self,
        vocab_size, embeded_size, pad_id,
        num_layers=1, drop_p=0,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embeded_size, padding_idx=pad_id)
        self.decoder = QRNN(embeded_size, embeded_size, n_layers=num_layers, dropout=drop_p, save_prev_x=True)
        self.classifier = nn.Linear(embeded_size, vocab_size)
        
    def forward(self, tgt_inp_ids, h):
        '''
            tgt_inp_ids: (bs, dec_seq_len)
            h: (num_decoder_layers, bs, embeded_size)
            returns: output, h
                output: (bs, dec_seq_len, dec_vocab_size)
                h: (num_decoder_layers, bs, embeded_size)
        '''
        embeded = self.embedding(tgt_inp_ids) # (bs, dec_seq_len, embeded_size)
        output, h = self.decoder(embeded, h) # (bs, dec_seq_len, embeded_size), (num_decoder_layers, bs, embeded_size)
        output = self.classifier(output) # (bs, dec_seq_len, dec_vocab_size)
        return output, h # (bs, dec_seq_len, dec_vocab_size), (num_decoder_layers, bs, embeded_size)
decoder = QRNNDecoder(dec_vocab_size, embeded_size, dec_pad_id, num_decoder_layers, drop_p)

tgt_input_ids = torch.randint(0, dec_vocab_size, (bs, dec_seq_len)) # (bs, dec_seq_len)
h = torch.randn((num_decoder_layers, bs, embeded_size)) # (num_decoder_layers, bs, embeded_size)

output, h = decoder(tgt_input_ids, h)
test_eq(output.shape, (bs, dec_seq_len, dec_vocab_size))
test_eq(h.shape, (num_decoder_layers, bs, embeded_size))<jupyter_output><empty_output><jupyter_text>CrossAttnQRNNDecoder For Decoder, uni-directional  CrossAttnQRNNLayer<jupyter_code># export
class CrossAttnQRNNLayer(nn.Module):
    def __init__(
        self, 
        targ_inp_size, embed_dim, window=1, # for qrnn_layer
        num_heads=1, drop_p=0, # for cross_attn
    ):
        super().__init__()
        self.qrnn_layer = QRNNLayer(targ_inp_size, hidden_size=embed_dim, save_prev_x=True, zoneout=0, window=window,
                                     output_gate=True, batch_first=True, backward=False)
        self.cross_attn = BatchFirstMultiheadAttention(embed_dim=embed_dim, num_heads=num_heads, dropout=drop_p)
    def forward(self, targ, memory, memory_key_padding_mask, hid=None):
        '''
            targ: (bs, targ_seq_len, targ_inp_size)
            memory:(bs, src_seq_len, embed_dim)
            memory_key_padding_mask: (bs, src_seq_len)
            hid: (bs, embed_dim)
            returns:
                output: (bs, tgt_seq_len, embed_dim)
                attn_weight: (bs, tgt_seq_len, src_seq_len)
                new_hid: (bs, targ_seq_len)
        '''
        qrnn_out, new_hid = self.qrnn_layer(targ, hid) # qrnn_out: (bs, targ_seq_len, embed_dim), new_hid: (bs, embed_dim)
        output, attn_weight = self.cross_attn(qrnn_out, memory, memory, key_padding_mask=memory_key_padding_mask) 
        return output, attn_weight, new_hid 
    def reset(self):
        self.qrnn_layer.reset()
targ = torch.randn((16, 18, 20))
memory = torch.randn((16, 28, 10))
memory_key_padding_mask = torch.zeros((16, 28)).bool()
cross_attn_qrnn_layer = CrossAttnQRNNLayer(20, 10)
output, attn_weight, new_hid = cross_attn_qrnn_layer(targ, memory, memory_key_padding_mask)
test_eq(output.shape, (16, 18, 10))
test_eq(attn_weight.shape, (16, 18, 28))
test_eq(new_hid.shape, (16, 10))<jupyter_output><empty_output><jupyter_text>CrossAttnQRNN<jupyter_code># export
class CrossAttnQRNN(nn.Module):
    def __init__(self, targ_inp_size, embed_dim, n_layers=1, num_heads=1, drop_p=0):
        super().__init__()
        self.layers = nn.ModuleList([CrossAttnQRNNLayer(targ_inp_size if l == 0 else embed_dim, embed_dim,
                                                       window=(2 if l ==0 else 1), num_heads=num_heads, drop_p=drop_p)
                                     for l in range(n_layers)])
        self.n_layers = n_layers
    def forward(self, targ, memory, memory_key_padding_mask, hid=None):
        '''
            targ: (bs, targ_seq_len, targ_inp_size)
            memory:(bs, src_seq_len, embed_dim)
            memory_key_padding_mask: (bs, src_seq_len)
            hid: (n_layers, bs, embed_dim)
            returns:
                output: (bs, tgt_seq_len, embed_dim)
                attn_weight: (n_layers, bs, tgt_seq_len, src_seq_len)
                hid: (n_layers, bs, targ_seq_len)
        '''
        assert hid.shape[0] == self.n_layers
        new_hid = []
        attn_weights = []
        for i, layer in enumerate(self.layers):
            targ, attn_weight, h = layer(targ, memory, memory_key_padding_mask, None if hid is None else hid[i])
            new_hid.append(h)
            attn_weights.append(attn_weight)
        return targ, torch.stack(attn_weights, 0), torch.stack(new_hid, 0)
    def reset(self):
        for layer in self.layers: layer.reset()
targ = torch.randn((16, 18, 20))
memory = torch.randn((16, 28, 10))
memory_key_padding_mask = torch.zeros((16, 28)).bool()
hid = torch.randn((2, 16, 10))
cross_attn_qrnn = CrossAttnQRNN(20, 10, 2)
output, attn_weight, new_hid = cross_attn_qrnn(targ, memory, memory_key_padding_mask, hid)
test_eq(output.shape, (16, 18, 10))
test_eq(attn_weight.shape, (2, 16, 18, 28))
test_eq(new_hid.shape, (2, 16, 10))
# export
class CrossAttnQRNNDecoder(nn.Module):
    def __init__(
        self,
        vocab_size, embeded_size, pad_id,
        num_layers=1, drop_p=0,
        num_heads=1,
    ):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embeded_size, padding_idx=pad_id)
        self.decoder = CrossAttnQRNN(embeded_size, embeded_size, n_layers=num_layers, num_heads=num_heads, drop_p=drop_p)
        self.classifier = nn.Linear(embeded_size, vocab_size)
        
    def forward(self, tgt_inp_ids, memory, memory_key_padding_mask, h):
        '''
            tgt_inp_ids: (bs, dec_seq_len)
            memory:(bs, src_seq_len, embed_dim)
            memory_key_padding_mask: (bs, src_seq_len)
            h: (num_decoder_layers, bs, embeded_size)
            returns: output, h
                output: (bs, dec_seq_len, dec_vocab_size)
                attn_weight: 
                h: (num_decoder_layers, bs, embeded_size)
        '''
        embeded = self.embedding(tgt_inp_ids) # (bs, dec_seq_len, embeded_size)
        output, attn_weight, h = self.decoder(embeded, memory, memory_key_padding_mask, h) # (bs, dec_seq_len, embeded_size), (num_decoder_layers, bs, embeded_size)
        output = self.classifier(output) # (bs, dec_seq_len, dec_vocab_size)
        return output, attn_weight, h # (bs, dec_seq_len, dec_vocab_size), (num_decoder_layers, bs, embeded_size)
tgt_inp_ids = torch.randint(0, 100, (16, 18))
memory = torch.randn((16, 28, 10))
memory_key_padding_mask = torch.zeros((16, 28)).bool()
hid = torch.randn((2, 16, 10))
cross_attn_qrnn_decoder = CrossAttnQRNNDecoder(100, 10, 1, num_layers=2)
output, attn_weight, new_hid = cross_attn_qrnn_decoder(tgt_inp_ids, memory, memory_key_padding_mask, hid)
test_eq(output.shape, (16, 18, 100))
test_eq(attn_weight.shape, (2, 16, 18, 28))
test_eq(new_hid.shape, (2, 16, 10))<jupyter_output><empty_output><jupyter_text>QRNN2QRNN<jupyter_code># export
class QRNN2QRNN(nn.Module):
    def __init__(
        self, 
        encoder: QRNNEncoder, decoder: QRNNDecoder, 
        num_encoder_layers, num_decoder_layers,
    ):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.proj = nn.Linear(2*num_encoder_layers, num_decoder_layers)
        
    def forward(self, src_input_ids, tgt_input_ids):
        '''
            src_input_ids: (bs, enc_seq_len)
            tgt_input_ids: (bs, dec_seq_len)
            returns: (bs, dec_seq_len, dec_vocab_size)
        '''
        _, enc_h = self.encoder(src_input_ids) # (2*num_encoder_layers, bs, embeded_size)
        enc_h = enc_h.permute(1, 2, 0) # (bs, embeded_size, 2*num_encoder_layers)
        enc_h_proj = self.proj(enc_h) # (bs, embeded_size, num_decoder_layers)
        enc_h_proj = enc_h_proj.permute(2, 0, 1) # (num_decoder_layers, bs, embeded_size)
        enc_h_proj = enc_h_proj.contiguous()
        
        output, dec_h = self.decoder(tgt_input_ids, enc_h_proj) # (bs, dec_seq_len, dec_vocab_size), (num_decoder_layers, bs, embeded_size)
        return output
qrnn2qrnn = QRNN2QRNN(encoder, decoder, num_encoder_layers, num_decoder_layers)

src_input_ids = torch.randint(0, enc_vocab_size, (bs, enc_seq_len)) # (bs, enc_seq_len)
tgt_input_ids = torch.randint(0, dec_vocab_size, (bs, dec_seq_len)) # (bs, dec_seq_len)
output = qrnn2qrnn(src_input_ids, tgt_input_ids) # (bs, dec_seq_len, dec_vocab_size)
test_eq(output.shape, (bs, dec_seq_len, dec_vocab_size))<jupyter_output><empty_output><jupyter_text>QRNN2AttnQRNN<jupyter_code># export
class QRNN2AttnQRNN(nn.Module):
    def __init__(
        self, 
        encoder: QRNNEncoder, decoder: CrossAttnQRNNDecoder, 
        num_encoder_layers, num_decoder_layers, embeded_size, # for project
        enc_pad_id,
    ):
        super().__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.hid_proj = nn.Linear(2*num_encoder_layers, num_decoder_layers)
        self.memory_proj = nn.Linear(2*embeded_size, embeded_size)
        
        self.enc_pad_id = enc_pad_id
        
    def forward(self, src_input_ids, tgt_input_ids):
        '''
            src_input_ids: (bs, enc_seq_len)
            tgt_input_ids: (bs, dec_seq_len)
            returns: (bs, dec_seq_len, dec_vocab_size)
        '''
        memory, enc_h = self.encoder(src_input_ids) # (bs, enc_seq_len, 2*embeded_size), (2*num_encoder_layers, bs, embeded_size)
        memory_proj = self.memory_proj(memory) # (bs, enc_seq_len, embeded_size)
        enc_h = enc_h.permute(1, 2, 0) # (bs, embeded_size, 2*num_encoder_layers)
        enc_h_proj = self.hid_proj(enc_h) # (bs, embeded_size, num_decoder_layers)
        enc_h_proj = enc_h_proj.permute(2, 0, 1) # (num_decoder_layers, bs, embeded_size)
        enc_h_proj = enc_h_proj.contiguous()
        memory_key_padding_mask = gen_key_padding_mask(src_input_ids, self.enc_pad_id)
        
        output, _, dec_h = self.decoder(tgt_input_ids, memory_proj, memory_key_padding_mask, enc_h_proj) # (bs, dec_seq_len, dec_vocab_size), (num_decoder_layers, bs, embeded_size)
        return output
cross_attn_qrnn_decoder = CrossAttnQRNNDecoder(dec_vocab_size, embeded_size, dec_pad_id, num_layers=2)
qrnn2attnqrnn = QRNN2AttnQRNN(encoder, cross_attn_qrnn_decoder, num_encoder_layers, num_decoder_layers, embeded_size, enc_pad_id)

src_input_ids = torch.randint(0, enc_vocab_size, (bs, enc_seq_len)) # (bs, enc_seq_len)
tgt_input_ids = torch.randint(0, dec_vocab_size, (bs, dec_seq_len)) # (bs, dec_seq_len)
output = qrnn2attnqrnn(src_input_ids, tgt_input_ids) # (bs, dec_seq_len, dec_vocab_size)
test_eq(output.shape, (bs, dec_seq_len, dec_vocab_size))<jupyter_output><empty_output><jupyter_text>GeneratedQRNN2QRNN<jupyter_code>enc_tokenizer = AutoTokenizer.from_pretrained('hfl/chinese-bert-wwm-ext')
dec_tokenizer = GPT2DecoderTokenizer.from_pretrained('distilgpt2')
# export
class GeneratedQRNN2QRNN():
    '''
        device is for created tensors
    '''
    def __init__(
        self,
        seq2seq: QRNN2QRNN, 
        enc_tokenizer: PreTrainedTokenizer,
        dec_tokenizer: PreTrainedTokenizer,
    ):
        self.seq2seq = seq2seq
        self.enc_tokenizer = enc_tokenizer
        self.dec_tokenizer = dec_tokenizer
        self.generatedLM = GeneratedLM(seq2seq.decoder, len(dec_tokenizer), dec_tokenizer.pad_token_id, [dec_tokenizer.eos_token_id], support_past=False)
# export
@patch
@torch.no_grad()
def generate_from_ids(self: GeneratedQRNN2QRNN, src_input_ids, generate_args: GenerateArgs):
    ''' src_input_ids: (bs, enc_seq_len) '''
    self.seq2seq.eval()
    self.seq2seq.decoder.decoder.reset()
    
    device = src_input_ids.device
    bs = src_input_ids.shape[0]
    tgt_input_ids = torch.zeros((bs, 1), dtype=torch.long).fill_(self.dec_tokenizer.bos_token_id).to(device) # (bs, 1)

    _, enc_h = self.seq2seq.encoder(src_input_ids) # (2*num_encoder_layers, bs, embeded_size)
    enc_h = enc_h.permute(1, 2, 0) # (bs, embeded_size, 2*num_encoder_layers)
    enc_h_proj =  self.seq2seq.proj(enc_h) # (bs, embeded_size, num_decoder_layers)
    model_otherargs = self.generatedLM.build_model_otherargs_for_beam([enc_h_proj], generate_args.num_beams) # (bs*num_beams, embeded_size, num_decoder_layers)
    enc_h_proj = model_otherargs[0].permute(2, 0, 1) # (num_decoder_layers, bs*num_beams, embeded_size)
    enc_h_proj = enc_h_proj.contiguous()

    result = self.generatedLM.generate(tgt_input_ids, generate_args, [enc_h_proj])

    return result
generated_qrnn2qrnn = GeneratedQRNN2QRNN(qrnn2qrnn, enc_tokenizer, dec_tokenizer)

generate_args = GenerateArgs(max_length=20, num_beams=2)
src_input_ids = torch.randint(0, enc_vocab_size, (bs, enc_seq_len)) # (bs, enc_seq_len)
result = generated_qrnn2qrnn.generate_from_ids(src_input_ids, generate_args)
test_eq(result.shape, (bs, 20))<jupyter_output><empty_output><jupyter_text>GeneratedQRNN2AttnQRNN<jupyter_code># export
class GeneratedQRNN2AttnQRNN():
    '''
        device is for created tensors
    '''
    def __init__(
        self,
        seq2seq: QRNN2AttnQRNN, 
        enc_tokenizer: PreTrainedTokenizer,
        dec_tokenizer: PreTrainedTokenizer,
    ):
        self.seq2seq = seq2seq
        self.enc_tokenizer = enc_tokenizer
        self.dec_tokenizer = dec_tokenizer
        self.generatedLM = GeneratedLM(seq2seq.decoder, len(dec_tokenizer), dec_tokenizer.pad_token_id, [dec_tokenizer.eos_token_id], support_past=False)
# export
@patch
@torch.no_grad()
def generate_from_ids(self: GeneratedQRNN2AttnQRNN, src_input_ids, generate_args: GenerateArgs):
    ''' src_input_ids: (bs, enc_seq_len) '''
    self.seq2seq.eval()
    self.seq2seq.decoder.decoder.reset()
    
    device = src_input_ids.device
    bs = src_input_ids.shape[0]
    tgt_input_ids = torch.zeros((bs, 1), dtype=torch.long).fill_(self.dec_tokenizer.bos_token_id).to(device) # (bs, 1)    
    
    memory, enc_h = self.seq2seq.encoder(src_input_ids) # (bs, enc_seq_len, 2*embeded_size), (2*num_encoder_layers, bs, embeded_size)
    memory_proj = self.seq2seq.memory_proj(memory) # (bs, enc_seq_len, embeded_size)
    memory_key_padding_mask = gen_key_padding_mask(src_input_ids, self.seq2seq.enc_pad_id)
    enc_h = enc_h.permute(1, 2, 0) # (bs, embeded_size, 2*num_encoder_layers)
    enc_h_proj = self.seq2seq.hid_proj(enc_h) # (bs, embeded_size, num_decoder_layers)
    memory_proj, memory_key_padding_mask, enc_h_proj = self.generatedLM.build_model_otherargs_for_beam([memory_proj, memory_key_padding_mask, enc_h_proj], generate_args.num_beams) # (bs*num_beams, embeded_size, num_decoder_layers)
    enc_h_proj = enc_h_proj.permute(2, 0, 1) # (num_decoder_layers, bs, embeded_size)
    enc_h_proj = enc_h_proj.contiguous()

    result = self.generatedLM.generate(tgt_input_ids, generate_args, [memory_proj, memory_key_padding_mask, enc_h_proj])

    return result
generated_qrnn2attnqrnn = GeneratedQRNN2AttnQRNN(qrnn2attnqrnn, enc_tokenizer, dec_tokenizer)

generate_args = GenerateArgs(max_length=20, num_beams=2)
src_input_ids = torch.randint(0, enc_vocab_size, (bs, enc_seq_len)) # (bs, enc_seq_len)
result = generated_qrnn2attnqrnn.generate_from_ids(src_input_ids, generate_args)
test_eq(result.shape, (bs, 20))<jupyter_output><empty_output><jupyter_text>Callback<jupyter_code># export
class QRNN2QRNNCallback(Callback):
    def after_batch(self):
        self.learn.model.decoder.decoder.reset()<jupyter_output><empty_output><jupyter_text>Export -<jupyter_code># hide
from nbdev.export import notebook2script
notebook2script()<jupyter_output>Converted 02_data.news_commentary.ipynb.
Converted 02_data.tatoeba.ipynb.
Converted 03a_models.patch.ipynb.
Converted 03c_models.bert2gpt2.ipynb.
Converted 03c_models.gru2gru.ipynb.
Converted 03c_models.qrnn2qrnn.ipynb.
Converted 03c_models.tran2tran.ipynb.
Converted 04_metrics.ipynb.
Converted 99_fulltest_bert2gpt2.ipynb.
Converted 99_fulltest_gru2gru.ipynb.
Converted 99_fulltest_qrnn2attnqrnn.ipynb.
Converted 99_fulltest_qrnn2qrnn.ipynb.
Converted 99_fulltest_tran2tran.ipynb.
Converted index.ipynb.
{'fraction_non_alphanumeric': 0.10900341335123943, 'fraction_numerical': 0.018409714061887975, 'mean_word_length': 4.025871766029247, 'pattern_counts': {'":': 0, '<': 44, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '11', 'dir': 'jupyter-structured-clean-dedup', 'id': '402736', 'license': 'Apache-2.0', 'n_tokens_mistral': 7235, 'n_tokens_neox': 6805, 'n_words': 1115, 'path': 'nbs/03c_models.qrnn2qrnn.ipynb', 'repo_name': 'cwza/nmt_try', 'size': '17871'}