from argparse import Namespace
from speech_translation.pretrain_schemes.asr_i6 import asr_i6

model_cfg = Namespace(
    encoder_ffn_embed_dim=2048,
    decoder_ffn_embed_dim=2048,
    encoder_embed_dim=512,
    decoder_output_dim=512,
    decoder_embed_dim=512,
    decoder_input_dim=512,
    encoder_attention_heads=8,
    decoder_attention_heads=8,
    dropout=0.1,
    attention_dropout=0.1,
    activation_dropout=0.1,
    embedding_dropout=0.1,
    encoder_layers=12,
    decoder_layers=12
)

criterion_cfg = Namespace(
    label_smoothing=0.1
)

for subepoch in range(1, 300):
    out = asr_i6(model_cfg, criterion_cfg, None, 0, subepoch)
    print('========')
    print(f'Subepoch {subepoch}')
    if out is not None:
        model, criterion = out
        print(model)
        print(criterion)
    else:
        print('Pretrain completed.')
    print('========')
