python - Trax のトランスフォーマーの導入例を理解する

Question

私の目標は、 https : //trax-ml.readthedocs.io/en/latest/notebooks/trax_intro.html にある Trax のトランスフォーマーの入門例を理解することです。

import trax

# Create a Transformer model.
# Pre-trained model config in gs://trax-ml/models/translation/ende_wmt32k.gin
model = trax.models.Transformer(
    input_vocab_size=33300,
    d_model=512, d_ff=2048,
    n_heads=8, n_encoder_layers=6, n_decoder_layers=6,
    max_len=2048, mode='predict')

# Initialize using pre-trained weights.
model.init_from_file('gs://trax-ml/models/translation/ende_wmt32k.pkl.gz',
                     weights_only=True)

# Tokenize a sentence.
sentence = 'It is nice to learn new things today!'
tokenized = list(trax.data.tokenize(iter([sentence]),  # Operates on streams.
                                    vocab_dir='gs://trax-ml/vocabs/',
                                    vocab_file='ende_32k.subword'))[0]

# Decode from the Transformer.
tokenized = tokenized[None, :]  # Add batch dimension.
tokenized_translation = trax.supervised.decoding.autoregressive_sample(
    model, tokenized, temperature=0.0)  # Higher temperature: more diverse results.

# De-tokenize,
tokenized_translation = tokenized_translation[0][:-1]  # Remove batch and EOS.
translation = trax.data.detokenize(tokenized_translation,
                                   vocab_dir='gs://trax-ml/vocabs/',
                                   vocab_file='ende_32k.subword')
print(translation)

この例はかなりうまく機能します。ただし、初期化されたモデルを使用して別の例を翻訳しようとすると、たとえば

sentence = 'I would like to try another example.'
tokenized = list(trax.data.tokenize(iter([sentence]),
                                    vocab_dir='gs://trax-ml/vocabs/',
                                    vocab_file='ende_32k.subword'))[0]
tokenized = tokenized[None, :]
tokenized_translation = trax.supervised.decoding.autoregressive_sample(
    model, tokenized, temperature=0.0)
tokenized_translation = tokenized_translation[0][:-1]
translation = trax.data.detokenize(tokenized_translation,
                                   vocab_dir='gs://trax-ml/vocabs/',
                                   vocab_file='ende_32k.subword')
print(translation)

!ローカル マシンと Google Colab で出力を取得します。他の例でも同じことが起こります。

新しいモデルを構築して初期化すると、すべて正常に動作します。

これはバグですか？そうでない場合、ここで何が起こっているのか、その動作を回避/修正するにはどうすればよいですか?

トークン化とトークン化解除はうまく機能しているようです。デバッグしました。で問題が発生したり、予期しないことが発生したようですtrax.supervised.decoding.autoregressive_sample。