{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [],
   "source": [
    "from fairseq.models.transformer import TransformerModel\n",
    "from fairseq import options, utils, checkpoint_utils\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "output_type": "stream",
     "name": "stdout",
     "text": "2020-10-06 18:37:24 | INFO | fairseq.file_utils | loading archive file ../../checkpoints/en2fr_wmt14_transformer/\n2020-10-06 18:37:32 | INFO | fairseq.tasks.translation | [en] dictionary: 44512 types\n2020-10-06 18:37:32 | INFO | fairseq.tasks.translation | [fr] dictionary: 44512 types\n2020-10-06 18:37:42 | INFO | fairseq.models.fairseq_model | Namespace(activation_dropout=0.0, activation_fn='relu', adam_betas='(0.9, 0.98)', adam_eps=1e-08, adaptive_input=False, adaptive_softmax_cutoff=None, adaptive_softmax_dropout=0, arch='transformer_vaswani_wmt_en_de_big', attention_dropout=0.0, bpe_codes='../../checkpoints/en2fr_wmt14_transformer/bpecodes', clip_norm=0.0, criterion='label_smoothed_cross_entropy', cross_self_attention=False, data='/data1$HOME/anaconda3/envs/fairenv/lib/python3.8/site-packages/fairseq/checkpoints/en2fr_wmt14_transformer', decoder_attention_heads=16, decoder_embed_dim=1024, decoder_embed_path=None, decoder_ffn_embed_dim=4096, decoder_input_dim=1024, decoder_layerdrop=0, decoder_layers=6, decoder_layers_to_keep=None, decoder_learned_pos=False, decoder_normalize_before=False, decoder_output_dim=1024, device_id=0, distributed_backend='nccl', distributed_init_method='tcp://learnfair0253:58342', distributed_port=58342, distributed_rank=0, distributed_world_size=128, dropout=0.1, encoder_attention_heads=16, encoder_embed_dim=1024, encoder_embed_path=None, encoder_ffn_embed_dim=4096, encoder_layerdrop=0, encoder_layers=6, encoder_layers_to_keep=None, encoder_learned_pos=False, encoder_normalize_before=False, eval_bleu_detok='space', eval_bleu_remove_bpe=None, eval_tokenized_bleu=False, fp16=True, label_smoothing=0.1, layernorm_embedding=False, left_pad_source=True, left_pad_target=False, log_format='json', log_interval=10, lr=[0.0007], lr_scheduler='inverse_sqrt', lr_shrink=0.1, max_epoch=0, max_sentences=None, max_sentences_valid=None, max_source_positions=1024, max_target_positions=1024, max_tokens=5120, max_update=80000, min_lr=1e-09, momentum=0.99, no_cross_attention=False, no_epoch_checkpoints=False, no_progress_bar=False, no_save=False, no_scale_embedding=False, no_token_positional_embeddings=False, num_batch_buckets=0, optimizer='adam', quant_noise_pq=0, quant_noise_pq_block_size=8, quant_noise_scalar=0, relu_dropout=0.0, restore_file='checkpoint_last.pt', sample_without_replacement=0, save_dir='/checkpoint02/myleott/2018-05-19/wmt14_en_fr.fp16_allreduce.fp16.maxupd80000.transformer_vaswani_wmt_en_de_big.shareemb.adam.beta0.9,0.98.initlr1e-07.warmup4000.lr0.0007.clip0.0.drop0.1.wd0.0.ls0.1.maxtok5120.seed2.ngpu128', save_interval=1, seed=2, sentence_avg=False, share_all_embeddings=True, share_decoder_input_output_embed=True, skip_invalid_size_inputs_valid_test=False, source_lang='en', target_lang='fr', task='translation', tie_adaptive_weights=False, train_subset='train', truncate_source=False, update_freq=[1.0], upsample_primary=1, use_old_adam=False, valid_subset='valid', validate_interval=1, warmup_init_lr=1e-07, warmup_updates=4000, weight_decay=0.0)\n"
    }
   ],
   "source": [
    "model = TransformerModel.from_pretrained(\n",
    "    '../../checkpoints/en2fr_wmt14_transformer/',\n",
    "    checkpoint_file='model.pt',\n",
    "    bpe_codes='../../checkpoints/en2fr_wmt14_transformer/bpecodes'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "'transformer_vaswani_wmt_en_de_big'"
     },
     "metadata": {},
     "execution_count": 57
    }
   ],
   "source": [
    "getattr(model.args, 'arch', 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "pretrain_state = checkpoint_utils.load_checkpoint_to_cpu(\n",
    "                '../../checkpoints/en2fr_wmt14_transformer/model.pt'\n",
    "            )\n",
    "state_dict = model.state_dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "260"
     },
     "metadata": {},
     "execution_count": 52
    }
   ],
   "source": [
    "len(state_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "184"
     },
     "metadata": {},
     "execution_count": 53
    }
   ],
   "source": [
    "len(pretrain_state['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "keys = state_dict.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_k = [k for k in keys if 'models.0' in k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "259"
     },
     "metadata": {},
     "execution_count": 46
    }
   ],
   "source": [
    "len(model_k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.3-final"
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}