{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5370b0ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "# nvcr.io/nvidia/tao/tao-toolkit-pyt:v3.22.05-py3\n",
    "import os\n",
    "import torch\n",
    "from core import checkpoint_encryption\n",
    "d = torch.load(\"/data/conformer.ckpt\", map_location=\"cpu\")\n",
    "\n",
    "# \"qahxrkjnrz\"\n",
    "# \"tlt_encode\"\n",
    "\n",
    "d = torch.load([\"/results/citrinet/train/checkpoints/\" + fn for fn in os.listdir(\"/results/citrinet/train/checkpoints/\") if \"ckpt\" in fn][0], map_location=\"cpu\")\n",
    "a = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")\n",
    "# import checkpoint_encryption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "43451d51",
   "metadata": {},
   "outputs": [],
   "source": [
    "# import torch\n",
    "\n",
    "# encryption_key = \"tlt_encode\"\n",
    "# encryption_key = \"wogsoimicu\"\n",
    "# # checkpoint_fp = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/encrypted_model_weights.ckpt\"\n",
    "# # checkpoint_fp = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/test.ckpt\"\n",
    "# checkpoint_fp = \"/mnt/data-ssd-1/tmp/tao/results/citrinet/train/checkpoints/trained-model--val_loss=129.1260-epoch=0.ckpt\"\n",
    "# state_dict = torch.load(checkpoint_fp, map_location=\"cpu\")\n",
    "\n",
    "# os.environ[\"HOST_DATA_DIR\"] = \"/mnt/data-ssd-1/tmp/tao/data\"\n",
    "# os.environ[\"HOST_SPECS_DIR\"] = \"/mnt/data-ssd-1/tmp/tao/specs\"\n",
    "# os.environ[\"HOST_RESULTS_DIR\"] = \"/mnt/data-ssd-1/tmp/tao/results\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9727ed01",
   "metadata": {},
   "outputs": [],
   "source": [
    "HOST_DATA_DIR=/mnt/data-ssd-1/tmp/tao/data\n",
    "HOST_SPECS_DIR=/mnt/data-ssd-1/tmp/tao/specs\n",
    "HOST_RESULTS_DIR=/mnt/data-ssd-1/tmp/tao/results\n",
    "\n",
    "DATA_DIR=/data\n",
    "SPECS_DIR=/specs\n",
    "RESULTS_DIR=/results\n",
    "KEY=tlt_encode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2abe216d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# speechtotext_en_us_conformer.tlt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b77556",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet download_specs \\\n",
    "    -r $RESULTS_DIR/speech_to_text_citrinet \\\n",
    "    -o $SPECS_DIR/speech_to_text_citrinet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af3836d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet dataset_convert \\\n",
    "    -e $SPECS_DIR/speech_to_text_citrinet/dataset_convert_an4.yaml \\\n",
    "    -r $RESULTS_DIR/citrinet/dataset_convert \\\n",
    "    source_data_dir=$DATA_DIR/an4 \\\n",
    "    target_data_dir=$DATA_DIR/an4_converted"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4ff3a17",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet create_tokenizer \\\n",
    "    -e $SPECS_DIR/speech_to_text_citrinet/create_tokenizer.yaml \\\n",
    "    -r $RESULTS_DIR/citrinet/create_tokenizer \\\n",
    "    manifests=$DATA_DIR/an4_converted/train_manifest.json \\\n",
    "    output_root=$DATA_DIR/an4 \\\n",
    "    vocab_size=32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7450c551",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet train \\\n",
    "     -e $SPECS_DIR/speech_to_text_citrinet/train_citrinet_bpe.yaml \\\n",
    "     -g 1 \\\n",
    "     -k tlt_encode \\\n",
    "     -r $RESULTS_DIR/citrinet/train \\\n",
    "     training_ds.manifest_filepath=$DATA_DIR/an4_converted/train_manifest.json \\\n",
    "     validation_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json \\\n",
    "     trainer.max_epochs=1 \\\n",
    "     training_ds.num_workers=4 \\\n",
    "     validation_ds.num_workers=4 \\\n",
    "     model.tokenizer.dir=$DATA_DIR/an4/tokenizer_spe_unigram_v32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "20f1f1fb",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d41794",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68e93d7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_conformer download_specs \\\n",
    "    -r $RESULTS_DIR/speech_to_text_conformer \\\n",
    "    -o $SPECS_DIR/speech_to_text_conformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9330c12",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_conformer finetune \\\n",
    "     -e $SPECS_DIR/speech_to_text_conformer/finetune_noop.yaml \\\n",
    "     -g 1 \\\n",
    "     -k $KEY \\\n",
    "     -m $RESULTS_DIR/conformer/train/checkpoints/speechtotext_en_us_conformer.tlt \\\n",
    "     -r $RESULTS_DIR/conformer/finetune_4 \\\n",
    "     finetuning_ds.manifest_filepath=$DATA_DIR/an4_converted/train_manifest.json \\\n",
    "     validation_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json \\\n",
    "     trainer.max_epochs=100 \\\n",
    "     finetuning_ds.num_workers=20 \\\n",
    "     validation_ds.num_workers=20 \\\n",
    "     trainer.gpus=1 \\\n",
    "     tokenizer.dir=$RESULTS_DIR/conformer/train/checkpoints/conformer-tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9bd15e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# tao speech_to_text_citrinet evaluate \\\n",
    "#      -e $SPECS_DIR/speech_to_text_conformer/evaluate.yaml \\\n",
    "#      -g 1 \\\n",
    "#      -k $KEY \\\n",
    "#      -m $RESULTS_DIR/conformer/train/checkpoints/speechtotext_en_us_conformer.tlt \\\n",
    "#      -r $RESULTS_DIR/conformer/evaluate \\\n",
    "#      test_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3cf1a867",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker exec -it d545f9 bash"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "62308717",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from core import checkpoint_encryption\n",
    "md = \"/results/conformer/finetune/checkpoints/\"\n",
    "d = torch.load([md + fn for fn in os.listdir(md) if \"ckpt\" in fn][0], map_location=\"cpu\")\n",
    "d2 = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")\n",
    "torch.save(d2[\"state_dict\"], \"/results/conformer/finetune/checkpoints/decrypted.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29cd65e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker cp be620aa63:/results/conformer/finetune_4/checkpoints/decrypted.pt ~/\n",
    "docker cp d545f9:/results/conformer/finetune_4/finetuned-model_epoch_11.tlt ~/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "72223554",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "d = torch.load(\"/home/georg/decrypted.pt\", map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "fe798767",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "odict_keys(['preprocessor.featurizer.window', 'preprocessor.featurizer.fb', 'encoder.pre_encode.out.weight', 'encoder.pre_encode.out.bias', 'encoder.pre_encode.conv.0.weight', 'encoder.pre_encode.conv.0.bias', 'encoder.pre_encode.conv.2.weight', 'encoder.pre_encode.conv.2.bias', 'encoder.layers.0.norm_feed_forward1.weight', 'encoder.layers.0.norm_feed_forward1.bias', 'encoder.layers.0.feed_forward1.linear1.weight', 'encoder.layers.0.feed_forward1.linear1.bias', 'encoder.layers.0.feed_forward1.linear2.weight', 'encoder.layers.0.feed_forward1.linear2.bias', 'encoder.layers.0.norm_conv.weight', 'encoder.layers.0.norm_conv.bias', 'encoder.layers.0.conv.pointwise_conv1.weight', 'encoder.layers.0.conv.pointwise_conv1.bias', 'encoder.layers.0.conv.depthwise_conv.weight', 'encoder.layers.0.conv.depthwise_conv.bias', 'encoder.layers.0.conv.batch_norm.weight', 'encoder.layers.0.conv.batch_norm.bias', 'encoder.layers.0.conv.batch_norm.running_mean', 'encoder.layers.0.conv.batch_norm.running_var', 'encoder.layers.0.conv.batch_norm.num_batches_tracked', 'encoder.layers.0.conv.pointwise_conv2.weight', 'encoder.layers.0.conv.pointwise_conv2.bias', 'encoder.layers.0.norm_self_att.weight', 'encoder.layers.0.norm_self_att.bias', 'encoder.layers.0.self_attn.pos_bias_u', 'encoder.layers.0.self_attn.pos_bias_v', 'encoder.layers.0.self_attn.linear_q.weight', 'encoder.layers.0.self_attn.linear_q.bias', 'encoder.layers.0.self_attn.linear_k.weight', 'encoder.layers.0.self_attn.linear_k.bias', 'encoder.layers.0.self_attn.linear_v.weight', 'encoder.layers.0.self_attn.linear_v.bias', 'encoder.layers.0.self_attn.linear_out.weight', 'encoder.layers.0.self_attn.linear_out.bias', 'encoder.layers.0.self_attn.linear_pos.weight', 'encoder.layers.0.norm_feed_forward2.weight', 'encoder.layers.0.norm_feed_forward2.bias', 'encoder.layers.0.feed_forward2.linear1.weight', 'encoder.layers.0.feed_forward2.linear1.bias', 'encoder.layers.0.feed_forward2.linear2.weight', 'encoder.layers.0.feed_forward2.linear2.bias', 'encoder.layers.0.norm_out.weight', 'encoder.layers.0.norm_out.bias', 'encoder.layers.1.norm_feed_forward1.weight', 'encoder.layers.1.norm_feed_forward1.bias', 'encoder.layers.1.feed_forward1.linear1.weight', 'encoder.layers.1.feed_forward1.linear1.bias', 'encoder.layers.1.feed_forward1.linear2.weight', 'encoder.layers.1.feed_forward1.linear2.bias', 'encoder.layers.1.norm_conv.weight', 'encoder.layers.1.norm_conv.bias', 'encoder.layers.1.conv.pointwise_conv1.weight', 'encoder.layers.1.conv.pointwise_conv1.bias', 'encoder.layers.1.conv.depthwise_conv.weight', 'encoder.layers.1.conv.depthwise_conv.bias', 'encoder.layers.1.conv.batch_norm.weight', 'encoder.layers.1.conv.batch_norm.bias', 'encoder.layers.1.conv.batch_norm.running_mean', 'encoder.layers.1.conv.batch_norm.running_var', 'encoder.layers.1.conv.batch_norm.num_batches_tracked', 'encoder.layers.1.conv.pointwise_conv2.weight', 'encoder.layers.1.conv.pointwise_conv2.bias', 'encoder.layers.1.norm_self_att.weight', 'encoder.layers.1.norm_self_att.bias', 'encoder.layers.1.self_attn.pos_bias_u', 'encoder.layers.1.self_attn.pos_bias_v', 'encoder.layers.1.self_attn.linear_q.weight', 'encoder.layers.1.self_attn.linear_q.bias', 'encoder.layers.1.self_attn.linear_k.weight', 'encoder.layers.1.self_attn.linear_k.bias', 'encoder.layers.1.self_attn.linear_v.weight', 'encoder.layers.1.self_attn.linear_v.bias', 'encoder.layers.1.self_attn.linear_out.weight', 'encoder.layers.1.self_attn.linear_out.bias', 'encoder.layers.1.self_attn.linear_pos.weight', 'encoder.layers.1.norm_feed_forward2.weight', 'encoder.layers.1.norm_feed_forward2.bias', 'encoder.layers.1.feed_forward2.linear1.weight', 'encoder.layers.1.feed_forward2.linear1.bias', 'encoder.layers.1.feed_forward2.linear2.weight', 'encoder.layers.1.feed_forward2.linear2.bias', 'encoder.layers.1.norm_out.weight', 'encoder.layers.1.norm_out.bias', 'encoder.layers.2.norm_feed_forward1.weight', 'encoder.layers.2.norm_feed_forward1.bias', 'encoder.layers.2.feed_forward1.linear1.weight', 'encoder.layers.2.feed_forward1.linear1.bias', 'encoder.layers.2.feed_forward1.linear2.weight', 'encoder.layers.2.feed_forward1.linear2.bias', 'encoder.layers.2.norm_conv.weight', 'encoder.layers.2.norm_conv.bias', 'encoder.layers.2.conv.pointwise_conv1.weight', 'encoder.layers.2.conv.pointwise_conv1.bias', 'encoder.layers.2.conv.depthwise_conv.weight', 'encoder.layers.2.conv.depthwise_conv.bias', 'encoder.layers.2.conv.batch_norm.weight', 'encoder.layers.2.conv.batch_norm.bias', 'encoder.layers.2.conv.batch_norm.running_mean', 'encoder.layers.2.conv.batch_norm.running_var', 'encoder.layers.2.conv.batch_norm.num_batches_tracked', 'encoder.layers.2.conv.pointwise_conv2.weight', 'encoder.layers.2.conv.pointwise_conv2.bias', 'encoder.layers.2.norm_self_att.weight', 'encoder.layers.2.norm_self_att.bias', 'encoder.layers.2.self_attn.pos_bias_u', 'encoder.layers.2.self_attn.pos_bias_v', 'encoder.layers.2.self_attn.linear_q.weight', 'encoder.layers.2.self_attn.linear_q.bias', 'encoder.layers.2.self_attn.linear_k.weight', 'encoder.layers.2.self_attn.linear_k.bias', 'encoder.layers.2.self_attn.linear_v.weight', 'encoder.layers.2.self_attn.linear_v.bias', 'encoder.layers.2.self_attn.linear_out.weight', 'encoder.layers.2.self_attn.linear_out.bias', 'encoder.layers.2.self_attn.linear_pos.weight', 'encoder.layers.2.norm_feed_forward2.weight', 'encoder.layers.2.norm_feed_forward2.bias', 'encoder.layers.2.feed_forward2.linear1.weight', 'encoder.layers.2.feed_forward2.linear1.bias', 'encoder.layers.2.feed_forward2.linear2.weight', 'encoder.layers.2.feed_forward2.linear2.bias', 'encoder.layers.2.norm_out.weight', 'encoder.layers.2.norm_out.bias', 'encoder.layers.3.norm_feed_forward1.weight', 'encoder.layers.3.norm_feed_forward1.bias', 'encoder.layers.3.feed_forward1.linear1.weight', 'encoder.layers.3.feed_forward1.linear1.bias', 'encoder.layers.3.feed_forward1.linear2.weight', 'encoder.layers.3.feed_forward1.linear2.bias', 'encoder.layers.3.norm_conv.weight', 'encoder.layers.3.norm_conv.bias', 'encoder.layers.3.conv.pointwise_conv1.weight', 'encoder.layers.3.conv.pointwise_conv1.bias', 'encoder.layers.3.conv.depthwise_conv.weight', 'encoder.layers.3.conv.depthwise_conv.bias', 'encoder.layers.3.conv.batch_norm.weight', 'encoder.layers.3.conv.batch_norm.bias', 'encoder.layers.3.conv.batch_norm.running_mean', 'encoder.layers.3.conv.batch_norm.running_var', 'encoder.layers.3.conv.batch_norm.num_batches_tracked', 'encoder.layers.3.conv.pointwise_conv2.weight', 'encoder.layers.3.conv.pointwise_conv2.bias', 'encoder.layers.3.norm_self_att.weight', 'encoder.layers.3.norm_self_att.bias', 'encoder.layers.3.self_attn.pos_bias_u', 'encoder.layers.3.self_attn.pos_bias_v', 'encoder.layers.3.self_attn.linear_q.weight', 'encoder.layers.3.self_attn.linear_q.bias', 'encoder.layers.3.self_attn.linear_k.weight', 'encoder.layers.3.self_attn.linear_k.bias', 'encoder.layers.3.self_attn.linear_v.weight', 'encoder.layers.3.self_attn.linear_v.bias', 'encoder.layers.3.self_attn.linear_out.weight', 'encoder.layers.3.self_attn.linear_out.bias', 'encoder.layers.3.self_attn.linear_pos.weight', 'encoder.layers.3.norm_feed_forward2.weight', 'encoder.layers.3.norm_feed_forward2.bias', 'encoder.layers.3.feed_forward2.linear1.weight', 'encoder.layers.3.feed_forward2.linear1.bias', 'encoder.layers.3.feed_forward2.linear2.weight', 'encoder.layers.3.feed_forward2.linear2.bias', 'encoder.layers.3.norm_out.weight', 'encoder.layers.3.norm_out.bias', 'encoder.layers.4.norm_feed_forward1.weight', 'encoder.layers.4.norm_feed_forward1.bias', 'encoder.layers.4.feed_forward1.linear1.weight', 'encoder.layers.4.feed_forward1.linear1.bias', 'encoder.layers.4.feed_forward1.linear2.weight', 'encoder.layers.4.feed_forward1.linear2.bias', 'encoder.layers.4.norm_conv.weight', 'encoder.layers.4.norm_conv.bias', 'encoder.layers.4.conv.pointwise_conv1.weight', 'encoder.layers.4.conv.pointwise_conv1.bias', 'encoder.layers.4.conv.depthwise_conv.weight', 'encoder.layers.4.conv.depthwise_conv.bias', 'encoder.layers.4.conv.batch_norm.weight', 'encoder.layers.4.conv.batch_norm.bias', 'encoder.layers.4.conv.batch_norm.running_mean', 'encoder.layers.4.conv.batch_norm.running_var', 'encoder.layers.4.conv.batch_norm.num_batches_tracked', 'encoder.layers.4.conv.pointwise_conv2.weight', 'encoder.layers.4.conv.pointwise_conv2.bias', 'encoder.layers.4.norm_self_att.weight', 'encoder.layers.4.norm_self_att.bias', 'encoder.layers.4.self_attn.pos_bias_u', 'encoder.layers.4.self_attn.pos_bias_v', 'encoder.layers.4.self_attn.linear_q.weight', 'encoder.layers.4.self_attn.linear_q.bias', 'encoder.layers.4.self_attn.linear_k.weight', 'encoder.layers.4.self_attn.linear_k.bias', 'encoder.layers.4.self_attn.linear_v.weight', 'encoder.layers.4.self_attn.linear_v.bias', 'encoder.layers.4.self_attn.linear_out.weight', 'encoder.layers.4.self_attn.linear_out.bias', 'encoder.layers.4.self_attn.linear_pos.weight', 'encoder.layers.4.norm_feed_forward2.weight', 'encoder.layers.4.norm_feed_forward2.bias', 'encoder.layers.4.feed_forward2.linear1.weight', 'encoder.layers.4.feed_forward2.linear1.bias', 'encoder.layers.4.feed_forward2.linear2.weight', 'encoder.layers.4.feed_forward2.linear2.bias', 'encoder.layers.4.norm_out.weight', 'encoder.layers.4.norm_out.bias', 'encoder.layers.5.norm_feed_forward1.weight', 'encoder.layers.5.norm_feed_forward1.bias', 'encoder.layers.5.feed_forward1.linear1.weight', 'encoder.layers.5.feed_forward1.linear1.bias', 'encoder.layers.5.feed_forward1.linear2.weight', 'encoder.layers.5.feed_forward1.linear2.bias', 'encoder.layers.5.norm_conv.weight', 'encoder.layers.5.norm_conv.bias', 'encoder.layers.5.conv.pointwise_conv1.weight', 'encoder.layers.5.conv.pointwise_conv1.bias', 'encoder.layers.5.conv.depthwise_conv.weight', 'encoder.layers.5.conv.depthwise_conv.bias', 'encoder.layers.5.conv.batch_norm.weight', 'encoder.layers.5.conv.batch_norm.bias', 'encoder.layers.5.conv.batch_norm.running_mean', 'encoder.layers.5.conv.batch_norm.running_var', 'encoder.layers.5.conv.batch_norm.num_batches_tracked', 'encoder.layers.5.conv.pointwise_conv2.weight', 'encoder.layers.5.conv.pointwise_conv2.bias', 'encoder.layers.5.norm_self_att.weight', 'encoder.layers.5.norm_self_att.bias', 'encoder.layers.5.self_attn.pos_bias_u', 'encoder.layers.5.self_attn.pos_bias_v', 'encoder.layers.5.self_attn.linear_q.weight', 'encoder.layers.5.self_attn.linear_q.bias', 'encoder.layers.5.self_attn.linear_k.weight', 'encoder.layers.5.self_attn.linear_k.bias', 'encoder.layers.5.self_attn.linear_v.weight', 'encoder.layers.5.self_attn.linear_v.bias', 'encoder.layers.5.self_attn.linear_out.weight', 'encoder.layers.5.self_attn.linear_out.bias', 'encoder.layers.5.self_attn.linear_pos.weight', 'encoder.layers.5.norm_feed_forward2.weight', 'encoder.layers.5.norm_feed_forward2.bias', 'encoder.layers.5.feed_forward2.linear1.weight', 'encoder.layers.5.feed_forward2.linear1.bias', 'encoder.layers.5.feed_forward2.linear2.weight', 'encoder.layers.5.feed_forward2.linear2.bias', 'encoder.layers.5.norm_out.weight', 'encoder.layers.5.norm_out.bias', 'encoder.layers.6.norm_feed_forward1.weight', 'encoder.layers.6.norm_feed_forward1.bias', 'encoder.layers.6.feed_forward1.linear1.weight', 'encoder.layers.6.feed_forward1.linear1.bias', 'encoder.layers.6.feed_forward1.linear2.weight', 'encoder.layers.6.feed_forward1.linear2.bias', 'encoder.layers.6.norm_conv.weight', 'encoder.layers.6.norm_conv.bias', 'encoder.layers.6.conv.pointwise_conv1.weight', 'encoder.layers.6.conv.pointwise_conv1.bias', 'encoder.layers.6.conv.depthwise_conv.weight', 'encoder.layers.6.conv.depthwise_conv.bias', 'encoder.layers.6.conv.batch_norm.weight', 'encoder.layers.6.conv.batch_norm.bias', 'encoder.layers.6.conv.batch_norm.running_mean', 'encoder.layers.6.conv.batch_norm.running_var', 'encoder.layers.6.conv.batch_norm.num_batches_tracked', 'encoder.layers.6.conv.pointwise_conv2.weight', 'encoder.layers.6.conv.pointwise_conv2.bias', 'encoder.layers.6.norm_self_att.weight', 'encoder.layers.6.norm_self_att.bias', 'encoder.layers.6.self_attn.pos_bias_u', 'encoder.layers.6.self_attn.pos_bias_v', 'encoder.layers.6.self_attn.linear_q.weight', 'encoder.layers.6.self_attn.linear_q.bias', 'encoder.layers.6.self_attn.linear_k.weight', 'encoder.layers.6.self_attn.linear_k.bias', 'encoder.layers.6.self_attn.linear_v.weight', 'encoder.layers.6.self_attn.linear_v.bias', 'encoder.layers.6.self_attn.linear_out.weight', 'encoder.layers.6.self_attn.linear_out.bias', 'encoder.layers.6.self_attn.linear_pos.weight', 'encoder.layers.6.norm_feed_forward2.weight', 'encoder.layers.6.norm_feed_forward2.bias', 'encoder.layers.6.feed_forward2.linear1.weight', 'encoder.layers.6.feed_forward2.linear1.bias', 'encoder.layers.6.feed_forward2.linear2.weight', 'encoder.layers.6.feed_forward2.linear2.bias', 'encoder.layers.6.norm_out.weight', 'encoder.layers.6.norm_out.bias', 'encoder.layers.7.norm_feed_forward1.weight', 'encoder.layers.7.norm_feed_forward1.bias', 'encoder.layers.7.feed_forward1.linear1.weight', 'encoder.layers.7.feed_forward1.linear1.bias', 'encoder.layers.7.feed_forward1.linear2.weight', 'encoder.layers.7.feed_forward1.linear2.bias', 'encoder.layers.7.norm_conv.weight', 'encoder.layers.7.norm_conv.bias', 'encoder.layers.7.conv.pointwise_conv1.weight', 'encoder.layers.7.conv.pointwise_conv1.bias', 'encoder.layers.7.conv.depthwise_conv.weight', 'encoder.layers.7.conv.depthwise_conv.bias', 'encoder.layers.7.conv.batch_norm.weight', 'encoder.layers.7.conv.batch_norm.bias', 'encoder.layers.7.conv.batch_norm.running_mean', 'encoder.layers.7.conv.batch_norm.running_var', 'encoder.layers.7.conv.batch_norm.num_batches_tracked', 'encoder.layers.7.conv.pointwise_conv2.weight', 'encoder.layers.7.conv.pointwise_conv2.bias', 'encoder.layers.7.norm_self_att.weight', 'encoder.layers.7.norm_self_att.bias', 'encoder.layers.7.self_attn.pos_bias_u', 'encoder.layers.7.self_attn.pos_bias_v', 'encoder.layers.7.self_attn.linear_q.weight', 'encoder.layers.7.self_attn.linear_q.bias', 'encoder.layers.7.self_attn.linear_k.weight', 'encoder.layers.7.self_attn.linear_k.bias', 'encoder.layers.7.self_attn.linear_v.weight', 'encoder.layers.7.self_attn.linear_v.bias', 'encoder.layers.7.self_attn.linear_out.weight', 'encoder.layers.7.self_attn.linear_out.bias', 'encoder.layers.7.self_attn.linear_pos.weight', 'encoder.layers.7.norm_feed_forward2.weight', 'encoder.layers.7.norm_feed_forward2.bias', 'encoder.layers.7.feed_forward2.linear1.weight', 'encoder.layers.7.feed_forward2.linear1.bias', 'encoder.layers.7.feed_forward2.linear2.weight', 'encoder.layers.7.feed_forward2.linear2.bias', 'encoder.layers.7.norm_out.weight', 'encoder.layers.7.norm_out.bias', 'encoder.layers.8.norm_feed_forward1.weight', 'encoder.layers.8.norm_feed_forward1.bias', 'encoder.layers.8.feed_forward1.linear1.weight', 'encoder.layers.8.feed_forward1.linear1.bias', 'encoder.layers.8.feed_forward1.linear2.weight', 'encoder.layers.8.feed_forward1.linear2.bias', 'encoder.layers.8.norm_conv.weight', 'encoder.layers.8.norm_conv.bias', 'encoder.layers.8.conv.pointwise_conv1.weight', 'encoder.layers.8.conv.pointwise_conv1.bias', 'encoder.layers.8.conv.depthwise_conv.weight', 'encoder.layers.8.conv.depthwise_conv.bias', 'encoder.layers.8.conv.batch_norm.weight', 'encoder.layers.8.conv.batch_norm.bias', 'encoder.layers.8.conv.batch_norm.running_mean', 'encoder.layers.8.conv.batch_norm.running_var', 'encoder.layers.8.conv.batch_norm.num_batches_tracked', 'encoder.layers.8.conv.pointwise_conv2.weight', 'encoder.layers.8.conv.pointwise_conv2.bias', 'encoder.layers.8.norm_self_att.weight', 'encoder.layers.8.norm_self_att.bias', 'encoder.layers.8.self_attn.pos_bias_u', 'encoder.layers.8.self_attn.pos_bias_v', 'encoder.layers.8.self_attn.linear_q.weight', 'encoder.layers.8.self_attn.linear_q.bias', 'encoder.layers.8.self_attn.linear_k.weight', 'encoder.layers.8.self_attn.linear_k.bias', 'encoder.layers.8.self_attn.linear_v.weight', 'encoder.layers.8.self_attn.linear_v.bias', 'encoder.layers.8.self_attn.linear_out.weight', 'encoder.layers.8.self_attn.linear_out.bias', 'encoder.layers.8.self_attn.linear_pos.weight', 'encoder.layers.8.norm_feed_forward2.weight', 'encoder.layers.8.norm_feed_forward2.bias', 'encoder.layers.8.feed_forward2.linear1.weight', 'encoder.layers.8.feed_forward2.linear1.bias', 'encoder.layers.8.feed_forward2.linear2.weight', 'encoder.layers.8.feed_forward2.linear2.bias', 'encoder.layers.8.norm_out.weight', 'encoder.layers.8.norm_out.bias', 'encoder.layers.9.norm_feed_forward1.weight', 'encoder.layers.9.norm_feed_forward1.bias', 'encoder.layers.9.feed_forward1.linear1.weight', 'encoder.layers.9.feed_forward1.linear1.bias', 'encoder.layers.9.feed_forward1.linear2.weight', 'encoder.layers.9.feed_forward1.linear2.bias', 'encoder.layers.9.norm_conv.weight', 'encoder.layers.9.norm_conv.bias', 'encoder.layers.9.conv.pointwise_conv1.weight', 'encoder.layers.9.conv.pointwise_conv1.bias', 'encoder.layers.9.conv.depthwise_conv.weight', 'encoder.layers.9.conv.depthwise_conv.bias', 'encoder.layers.9.conv.batch_norm.weight', 'encoder.layers.9.conv.batch_norm.bias', 'encoder.layers.9.conv.batch_norm.running_mean', 'encoder.layers.9.conv.batch_norm.running_var', 'encoder.layers.9.conv.batch_norm.num_batches_tracked', 'encoder.layers.9.conv.pointwise_conv2.weight', 'encoder.layers.9.conv.pointwise_conv2.bias', 'encoder.layers.9.norm_self_att.weight', 'encoder.layers.9.norm_self_att.bias', 'encoder.layers.9.self_attn.pos_bias_u', 'encoder.layers.9.self_attn.pos_bias_v', 'encoder.layers.9.self_attn.linear_q.weight', 'encoder.layers.9.self_attn.linear_q.bias', 'encoder.layers.9.self_attn.linear_k.weight', 'encoder.layers.9.self_attn.linear_k.bias', 'encoder.layers.9.self_attn.linear_v.weight', 'encoder.layers.9.self_attn.linear_v.bias', 'encoder.layers.9.self_attn.linear_out.weight', 'encoder.layers.9.self_attn.linear_out.bias', 'encoder.layers.9.self_attn.linear_pos.weight', 'encoder.layers.9.norm_feed_forward2.weight', 'encoder.layers.9.norm_feed_forward2.bias', 'encoder.layers.9.feed_forward2.linear1.weight', 'encoder.layers.9.feed_forward2.linear1.bias', 'encoder.layers.9.feed_forward2.linear2.weight', 'encoder.layers.9.feed_forward2.linear2.bias', 'encoder.layers.9.norm_out.weight', 'encoder.layers.9.norm_out.bias', 'encoder.layers.10.norm_feed_forward1.weight', 'encoder.layers.10.norm_feed_forward1.bias', 'encoder.layers.10.feed_forward1.linear1.weight', 'encoder.layers.10.feed_forward1.linear1.bias', 'encoder.layers.10.feed_forward1.linear2.weight', 'encoder.layers.10.feed_forward1.linear2.bias', 'encoder.layers.10.norm_conv.weight', 'encoder.layers.10.norm_conv.bias', 'encoder.layers.10.conv.pointwise_conv1.weight', 'encoder.layers.10.conv.pointwise_conv1.bias', 'encoder.layers.10.conv.depthwise_conv.weight', 'encoder.layers.10.conv.depthwise_conv.bias', 'encoder.layers.10.conv.batch_norm.weight', 'encoder.layers.10.conv.batch_norm.bias', 'encoder.layers.10.conv.batch_norm.running_mean', 'encoder.layers.10.conv.batch_norm.running_var', 'encoder.layers.10.conv.batch_norm.num_batches_tracked', 'encoder.layers.10.conv.pointwise_conv2.weight', 'encoder.layers.10.conv.pointwise_conv2.bias', 'encoder.layers.10.norm_self_att.weight', 'encoder.layers.10.norm_self_att.bias', 'encoder.layers.10.self_attn.pos_bias_u', 'encoder.layers.10.self_attn.pos_bias_v', 'encoder.layers.10.self_attn.linear_q.weight', 'encoder.layers.10.self_attn.linear_q.bias', 'encoder.layers.10.self_attn.linear_k.weight', 'encoder.layers.10.self_attn.linear_k.bias', 'encoder.layers.10.self_attn.linear_v.weight', 'encoder.layers.10.self_attn.linear_v.bias', 'encoder.layers.10.self_attn.linear_out.weight', 'encoder.layers.10.self_attn.linear_out.bias', 'encoder.layers.10.self_attn.linear_pos.weight', 'encoder.layers.10.norm_feed_forward2.weight', 'encoder.layers.10.norm_feed_forward2.bias', 'encoder.layers.10.feed_forward2.linear1.weight', 'encoder.layers.10.feed_forward2.linear1.bias', 'encoder.layers.10.feed_forward2.linear2.weight', 'encoder.layers.10.feed_forward2.linear2.bias', 'encoder.layers.10.norm_out.weight', 'encoder.layers.10.norm_out.bias', 'encoder.layers.11.norm_feed_forward1.weight', 'encoder.layers.11.norm_feed_forward1.bias', 'encoder.layers.11.feed_forward1.linear1.weight', 'encoder.layers.11.feed_forward1.linear1.bias', 'encoder.layers.11.feed_forward1.linear2.weight', 'encoder.layers.11.feed_forward1.linear2.bias', 'encoder.layers.11.norm_conv.weight', 'encoder.layers.11.norm_conv.bias', 'encoder.layers.11.conv.pointwise_conv1.weight', 'encoder.layers.11.conv.pointwise_conv1.bias', 'encoder.layers.11.conv.depthwise_conv.weight', 'encoder.layers.11.conv.depthwise_conv.bias', 'encoder.layers.11.conv.batch_norm.weight', 'encoder.layers.11.conv.batch_norm.bias', 'encoder.layers.11.conv.batch_norm.running_mean', 'encoder.layers.11.conv.batch_norm.running_var', 'encoder.layers.11.conv.batch_norm.num_batches_tracked', 'encoder.layers.11.conv.pointwise_conv2.weight', 'encoder.layers.11.conv.pointwise_conv2.bias', 'encoder.layers.11.norm_self_att.weight', 'encoder.layers.11.norm_self_att.bias', 'encoder.layers.11.self_attn.pos_bias_u', 'encoder.layers.11.self_attn.pos_bias_v', 'encoder.layers.11.self_attn.linear_q.weight', 'encoder.layers.11.self_attn.linear_q.bias', 'encoder.layers.11.self_attn.linear_k.weight', 'encoder.layers.11.self_attn.linear_k.bias', 'encoder.layers.11.self_attn.linear_v.weight', 'encoder.layers.11.self_attn.linear_v.bias', 'encoder.layers.11.self_attn.linear_out.weight', 'encoder.layers.11.self_attn.linear_out.bias', 'encoder.layers.11.self_attn.linear_pos.weight', 'encoder.layers.11.norm_feed_forward2.weight', 'encoder.layers.11.norm_feed_forward2.bias', 'encoder.layers.11.feed_forward2.linear1.weight', 'encoder.layers.11.feed_forward2.linear1.bias', 'encoder.layers.11.feed_forward2.linear2.weight', 'encoder.layers.11.feed_forward2.linear2.bias', 'encoder.layers.11.norm_out.weight', 'encoder.layers.11.norm_out.bias', 'encoder.layers.12.norm_feed_forward1.weight', 'encoder.layers.12.norm_feed_forward1.bias', 'encoder.layers.12.feed_forward1.linear1.weight', 'encoder.layers.12.feed_forward1.linear1.bias', 'encoder.layers.12.feed_forward1.linear2.weight', 'encoder.layers.12.feed_forward1.linear2.bias', 'encoder.layers.12.norm_conv.weight', 'encoder.layers.12.norm_conv.bias', 'encoder.layers.12.conv.pointwise_conv1.weight', 'encoder.layers.12.conv.pointwise_conv1.bias', 'encoder.layers.12.conv.depthwise_conv.weight', 'encoder.layers.12.conv.depthwise_conv.bias', 'encoder.layers.12.conv.batch_norm.weight', 'encoder.layers.12.conv.batch_norm.bias', 'encoder.layers.12.conv.batch_norm.running_mean', 'encoder.layers.12.conv.batch_norm.running_var', 'encoder.layers.12.conv.batch_norm.num_batches_tracked', 'encoder.layers.12.conv.pointwise_conv2.weight', 'encoder.layers.12.conv.pointwise_conv2.bias', 'encoder.layers.12.norm_self_att.weight', 'encoder.layers.12.norm_self_att.bias', 'encoder.layers.12.self_attn.pos_bias_u', 'encoder.layers.12.self_attn.pos_bias_v', 'encoder.layers.12.self_attn.linear_q.weight', 'encoder.layers.12.self_attn.linear_q.bias', 'encoder.layers.12.self_attn.linear_k.weight', 'encoder.layers.12.self_attn.linear_k.bias', 'encoder.layers.12.self_attn.linear_v.weight', 'encoder.layers.12.self_attn.linear_v.bias', 'encoder.layers.12.self_attn.linear_out.weight', 'encoder.layers.12.self_attn.linear_out.bias', 'encoder.layers.12.self_attn.linear_pos.weight', 'encoder.layers.12.norm_feed_forward2.weight', 'encoder.layers.12.norm_feed_forward2.bias', 'encoder.layers.12.feed_forward2.linear1.weight', 'encoder.layers.12.feed_forward2.linear1.bias', 'encoder.layers.12.feed_forward2.linear2.weight', 'encoder.layers.12.feed_forward2.linear2.bias', 'encoder.layers.12.norm_out.weight', 'encoder.layers.12.norm_out.bias', 'encoder.layers.13.norm_feed_forward1.weight', 'encoder.layers.13.norm_feed_forward1.bias', 'encoder.layers.13.feed_forward1.linear1.weight', 'encoder.layers.13.feed_forward1.linear1.bias', 'encoder.layers.13.feed_forward1.linear2.weight', 'encoder.layers.13.feed_forward1.linear2.bias', 'encoder.layers.13.norm_conv.weight', 'encoder.layers.13.norm_conv.bias', 'encoder.layers.13.conv.pointwise_conv1.weight', 'encoder.layers.13.conv.pointwise_conv1.bias', 'encoder.layers.13.conv.depthwise_conv.weight', 'encoder.layers.13.conv.depthwise_conv.bias', 'encoder.layers.13.conv.batch_norm.weight', 'encoder.layers.13.conv.batch_norm.bias', 'encoder.layers.13.conv.batch_norm.running_mean', 'encoder.layers.13.conv.batch_norm.running_var', 'encoder.layers.13.conv.batch_norm.num_batches_tracked', 'encoder.layers.13.conv.pointwise_conv2.weight', 'encoder.layers.13.conv.pointwise_conv2.bias', 'encoder.layers.13.norm_self_att.weight', 'encoder.layers.13.norm_self_att.bias', 'encoder.layers.13.self_attn.pos_bias_u', 'encoder.layers.13.self_attn.pos_bias_v', 'encoder.layers.13.self_attn.linear_q.weight', 'encoder.layers.13.self_attn.linear_q.bias', 'encoder.layers.13.self_attn.linear_k.weight', 'encoder.layers.13.self_attn.linear_k.bias', 'encoder.layers.13.self_attn.linear_v.weight', 'encoder.layers.13.self_attn.linear_v.bias', 'encoder.layers.13.self_attn.linear_out.weight', 'encoder.layers.13.self_attn.linear_out.bias', 'encoder.layers.13.self_attn.linear_pos.weight', 'encoder.layers.13.norm_feed_forward2.weight', 'encoder.layers.13.norm_feed_forward2.bias', 'encoder.layers.13.feed_forward2.linear1.weight', 'encoder.layers.13.feed_forward2.linear1.bias', 'encoder.layers.13.feed_forward2.linear2.weight', 'encoder.layers.13.feed_forward2.linear2.bias', 'encoder.layers.13.norm_out.weight', 'encoder.layers.13.norm_out.bias', 'encoder.layers.14.norm_feed_forward1.weight', 'encoder.layers.14.norm_feed_forward1.bias', 'encoder.layers.14.feed_forward1.linear1.weight', 'encoder.layers.14.feed_forward1.linear1.bias', 'encoder.layers.14.feed_forward1.linear2.weight', 'encoder.layers.14.feed_forward1.linear2.bias', 'encoder.layers.14.norm_conv.weight', 'encoder.layers.14.norm_conv.bias', 'encoder.layers.14.conv.pointwise_conv1.weight', 'encoder.layers.14.conv.pointwise_conv1.bias', 'encoder.layers.14.conv.depthwise_conv.weight', 'encoder.layers.14.conv.depthwise_conv.bias', 'encoder.layers.14.conv.batch_norm.weight', 'encoder.layers.14.conv.batch_norm.bias', 'encoder.layers.14.conv.batch_norm.running_mean', 'encoder.layers.14.conv.batch_norm.running_var', 'encoder.layers.14.conv.batch_norm.num_batches_tracked', 'encoder.layers.14.conv.pointwise_conv2.weight', 'encoder.layers.14.conv.pointwise_conv2.bias', 'encoder.layers.14.norm_self_att.weight', 'encoder.layers.14.norm_self_att.bias', 'encoder.layers.14.self_attn.pos_bias_u', 'encoder.layers.14.self_attn.pos_bias_v', 'encoder.layers.14.self_attn.linear_q.weight', 'encoder.layers.14.self_attn.linear_q.bias', 'encoder.layers.14.self_attn.linear_k.weight', 'encoder.layers.14.self_attn.linear_k.bias', 'encoder.layers.14.self_attn.linear_v.weight', 'encoder.layers.14.self_attn.linear_v.bias', 'encoder.layers.14.self_attn.linear_out.weight', 'encoder.layers.14.self_attn.linear_out.bias', 'encoder.layers.14.self_attn.linear_pos.weight', 'encoder.layers.14.norm_feed_forward2.weight', 'encoder.layers.14.norm_feed_forward2.bias', 'encoder.layers.14.feed_forward2.linear1.weight', 'encoder.layers.14.feed_forward2.linear1.bias', 'encoder.layers.14.feed_forward2.linear2.weight', 'encoder.layers.14.feed_forward2.linear2.bias', 'encoder.layers.14.norm_out.weight', 'encoder.layers.14.norm_out.bias', 'encoder.layers.15.norm_feed_forward1.weight', 'encoder.layers.15.norm_feed_forward1.bias', 'encoder.layers.15.feed_forward1.linear1.weight', 'encoder.layers.15.feed_forward1.linear1.bias', 'encoder.layers.15.feed_forward1.linear2.weight', 'encoder.layers.15.feed_forward1.linear2.bias', 'encoder.layers.15.norm_conv.weight', 'encoder.layers.15.norm_conv.bias', 'encoder.layers.15.conv.pointwise_conv1.weight', 'encoder.layers.15.conv.pointwise_conv1.bias', 'encoder.layers.15.conv.depthwise_conv.weight', 'encoder.layers.15.conv.depthwise_conv.bias', 'encoder.layers.15.conv.batch_norm.weight', 'encoder.layers.15.conv.batch_norm.bias', 'encoder.layers.15.conv.batch_norm.running_mean', 'encoder.layers.15.conv.batch_norm.running_var', 'encoder.layers.15.conv.batch_norm.num_batches_tracked', 'encoder.layers.15.conv.pointwise_conv2.weight', 'encoder.layers.15.conv.pointwise_conv2.bias', 'encoder.layers.15.norm_self_att.weight', 'encoder.layers.15.norm_self_att.bias', 'encoder.layers.15.self_attn.pos_bias_u', 'encoder.layers.15.self_attn.pos_bias_v', 'encoder.layers.15.self_attn.linear_q.weight', 'encoder.layers.15.self_attn.linear_q.bias', 'encoder.layers.15.self_attn.linear_k.weight', 'encoder.layers.15.self_attn.linear_k.bias', 'encoder.layers.15.self_attn.linear_v.weight', 'encoder.layers.15.self_attn.linear_v.bias', 'encoder.layers.15.self_attn.linear_out.weight', 'encoder.layers.15.self_attn.linear_out.bias', 'encoder.layers.15.self_attn.linear_pos.weight', 'encoder.layers.15.norm_feed_forward2.weight', 'encoder.layers.15.norm_feed_forward2.bias', 'encoder.layers.15.feed_forward2.linear1.weight', 'encoder.layers.15.feed_forward2.linear1.bias', 'encoder.layers.15.feed_forward2.linear2.weight', 'encoder.layers.15.feed_forward2.linear2.bias', 'encoder.layers.15.norm_out.weight', 'encoder.layers.15.norm_out.bias', 'encoder.layers.16.norm_feed_forward1.weight', 'encoder.layers.16.norm_feed_forward1.bias', 'encoder.layers.16.feed_forward1.linear1.weight', 'encoder.layers.16.feed_forward1.linear1.bias', 'encoder.layers.16.feed_forward1.linear2.weight', 'encoder.layers.16.feed_forward1.linear2.bias', 'encoder.layers.16.norm_conv.weight', 'encoder.layers.16.norm_conv.bias', 'encoder.layers.16.conv.pointwise_conv1.weight', 'encoder.layers.16.conv.pointwise_conv1.bias', 'encoder.layers.16.conv.depthwise_conv.weight', 'encoder.layers.16.conv.depthwise_conv.bias', 'encoder.layers.16.conv.batch_norm.weight', 'encoder.layers.16.conv.batch_norm.bias', 'encoder.layers.16.conv.batch_norm.running_mean', 'encoder.layers.16.conv.batch_norm.running_var', 'encoder.layers.16.conv.batch_norm.num_batches_tracked', 'encoder.layers.16.conv.pointwise_conv2.weight', 'encoder.layers.16.conv.pointwise_conv2.bias', 'encoder.layers.16.norm_self_att.weight', 'encoder.layers.16.norm_self_att.bias', 'encoder.layers.16.self_attn.pos_bias_u', 'encoder.layers.16.self_attn.pos_bias_v', 'encoder.layers.16.self_attn.linear_q.weight', 'encoder.layers.16.self_attn.linear_q.bias', 'encoder.layers.16.self_attn.linear_k.weight', 'encoder.layers.16.self_attn.linear_k.bias', 'encoder.layers.16.self_attn.linear_v.weight', 'encoder.layers.16.self_attn.linear_v.bias', 'encoder.layers.16.self_attn.linear_out.weight', 'encoder.layers.16.self_attn.linear_out.bias', 'encoder.layers.16.self_attn.linear_pos.weight', 'encoder.layers.16.norm_feed_forward2.weight', 'encoder.layers.16.norm_feed_forward2.bias', 'encoder.layers.16.feed_forward2.linear1.weight', 'encoder.layers.16.feed_forward2.linear1.bias', 'encoder.layers.16.feed_forward2.linear2.weight', 'encoder.layers.16.feed_forward2.linear2.bias', 'encoder.layers.16.norm_out.weight', 'encoder.layers.16.norm_out.bias', 'encoder.layers.17.norm_feed_forward1.weight', 'encoder.layers.17.norm_feed_forward1.bias', 'encoder.layers.17.feed_forward1.linear1.weight', 'encoder.layers.17.feed_forward1.linear1.bias', 'encoder.layers.17.feed_forward1.linear2.weight', 'encoder.layers.17.feed_forward1.linear2.bias', 'encoder.layers.17.norm_conv.weight', 'encoder.layers.17.norm_conv.bias', 'encoder.layers.17.conv.pointwise_conv1.weight', 'encoder.layers.17.conv.pointwise_conv1.bias', 'encoder.layers.17.conv.depthwise_conv.weight', 'encoder.layers.17.conv.depthwise_conv.bias', 'encoder.layers.17.conv.batch_norm.weight', 'encoder.layers.17.conv.batch_norm.bias', 'encoder.layers.17.conv.batch_norm.running_mean', 'encoder.layers.17.conv.batch_norm.running_var', 'encoder.layers.17.conv.batch_norm.num_batches_tracked', 'encoder.layers.17.conv.pointwise_conv2.weight', 'encoder.layers.17.conv.pointwise_conv2.bias', 'encoder.layers.17.norm_self_att.weight', 'encoder.layers.17.norm_self_att.bias', 'encoder.layers.17.self_attn.pos_bias_u', 'encoder.layers.17.self_attn.pos_bias_v', 'encoder.layers.17.self_attn.linear_q.weight', 'encoder.layers.17.self_attn.linear_q.bias', 'encoder.layers.17.self_attn.linear_k.weight', 'encoder.layers.17.self_attn.linear_k.bias', 'encoder.layers.17.self_attn.linear_v.weight', 'encoder.layers.17.self_attn.linear_v.bias', 'encoder.layers.17.self_attn.linear_out.weight', 'encoder.layers.17.self_attn.linear_out.bias', 'encoder.layers.17.self_attn.linear_pos.weight', 'encoder.layers.17.norm_feed_forward2.weight', 'encoder.layers.17.norm_feed_forward2.bias', 'encoder.layers.17.feed_forward2.linear1.weight', 'encoder.layers.17.feed_forward2.linear1.bias', 'encoder.layers.17.feed_forward2.linear2.weight', 'encoder.layers.17.feed_forward2.linear2.bias', 'encoder.layers.17.norm_out.weight', 'encoder.layers.17.norm_out.bias', 'decoder.decoder_layers.0.weight', 'decoder.decoder_layers.0.bias'])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d.keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "9cead1ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from eff. import checkpoint_encryption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "f75801f8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/home/georg/notebooks/tasks/s2t\r\n"
     ]
    }
   ],
   "source": [
    "!pwd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "618b4f82",
   "metadata": {},
   "outputs": [],
   "source": [
    "import eff\n",
    "\n",
    "\n",
    "passphrase = \"tlt_encode\"\n",
    "encryption = \"qahxrkjnrz\"\n",
    "\n",
    "fp_in = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/encrypted_model_weights.ckpt\"\n",
    "fp_out = \"/home/georg/notebooks/tasks/s2t/decrypted.ckpt\"\n",
    "\n",
    "# eff.decrypt_file(fp_in, fp_out, key_1, \"blba\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "id": "d62177a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "eff.decrypt_file(fp_in, fp_out, key_2, key_1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b44799f",
   "metadata": {},
   "outputs": [],
   "source": [
    "eff.decrypt_stream"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e64a7a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import eff\n",
    "\n",
    "key_1 = \"qahxrkjnrz\"\n",
    "key_2 = \"tlt_encode\"\n",
    "\n",
    "fp_in = \"/home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/encrypted_model_weights.ckpt\"\n",
    "fp_out = \"/home/georg/notebooks/tasks/s2t/decrypted.ckpt\"\n",
    "\n",
    "# eff.decrypt_file(fp_in, fp_out, key_1, \"blba\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "efe2abc3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 87,
   "id": "5cbe449d",
   "metadata": {},
   "outputs": [],
   "source": [
    "eff.codec.decrypt_file(fp_in, fp_out, key_1.encode(), \"Random\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 90,
   "id": "a8f43f07",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "bad magic number\r\n"
     ]
    }
   ],
   "source": [
    "!openssl enc -in /home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/encrypted_model_weights.ckpt -out /home/georg/notebooks/tasks/s2t/decrypted.ckpt -d -aes-256-ctr -pbkdf2 -k qahxrkjnrz"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "id": "7d56e373",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'kkqfyajsee'"
      ]
     },
     "execution_count": 91,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eff.codec.get_random_encryption()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e70893bc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c79af2fc",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "ad299d2b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "f56a76ec19b806a8f19371900bf287592a51e672f3d04e156db480d1d3b6e514  /home/georg/notebooks/tasks/s2t/decrypted.ckpt\r\n"
     ]
    }
   ],
   "source": [
    "# f56a76ec19b806a8f19371900bf287592a51e672f3d04e156db480d1d3b6e514\n",
    "!sha256sum /home/georg/notebooks/tasks/s2t/decrypted.ckpt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f38e6e36",
   "metadata": {},
   "outputs": [],
   "source": [
    "256 PBKDF2 AES CTR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b44c6d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "openssl aes-256-cbc -pbkdf2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 88,
   "id": "b971633e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "b'qahxrkjnrz'"
      ]
     },
     "execution_count": 88,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\"qahxrkjnrz\".encode(\"utf8\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32def966",
   "metadata": {},
   "outputs": [],
   "source": [
    "Derives the 256 bit key using PBKDF2 passphrase stretching and decodes\n",
    "a stream encoded using AES in CTR mode by the above encode function.\n",
    "Processes 1024 bytes at a time and uses the 'nonce' value included at\n",
    "the beginning of the cipher text input stream\n",
    "\n",
    "passphrase (bytes): byte text representing the phrase used during generation of encryption key.\n",
    "encryption: (string | bool) identifier of used encryption (DEFAULT: \"Random\")."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "id": "939b045e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "d = torch.load(fp_out, map_location=\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aefb6d31",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfc12825",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff6900f9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b97743c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ba41923",
   "metadata": {},
   "outputs": [],
   "source": [
    "# nvcr.io/nvidia/tao/tao-toolkit-pyt:v3.22.05-py3\n",
    "import os\n",
    "import torch\n",
    "from core import checkpoint_encryption\n",
    "d = torch.load(\"/data/conformer.ckpt\", map_location=\"cpu\")\n",
    "\n",
    "# \"qahxrkjnrz\"\n",
    "# \"tlt_encode\"\n",
    "\n",
    "d = torch.load([\"/results/citrinet/train/checkpoints/\" + fn for fn in os.listdir(\"/results/citrinet/train/checkpoints/\") if \"ckpt\" in fn][0], map_location=\"cpu\")\n",
    "a = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")\n",
    "# import checkpoint_encryption"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97d2891a",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2760b786",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e93a177",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83d444ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "[NeMo I 2022-09-12 23:20:02 finetune:128] Model restored from '/results/conformer/train/checkpoints/speechtotext_en_us_conformer.tlt'\n",
    "[NeMo W 2022-09-12 23:20:02 modelPT:215] You tried to register an artifact under config key=tokenizer.model_path but an artifact for it has already been registered.\n",
    "[NeMo W 2022-09-12 23:20:02 modelPT:215] You tried to register an artifact under config key=tokenizer.vocab_path but an artifact for it has already been registered.\n",
    "[NeMo W 2022-09-12 23:20:02 modelPT:215] You tried to register an artifact under config key=tokenizer.spe_tokenizer_vocab but an artifact for it has already been registered.\n",
    "[NeMo I 2022-09-12 23:20:02 mixins:146] Tokenizer SentencePieceTokenizer initialized with 256 tokens\n",
    "[NeMo I 2022-09-12 23:20:03 ctc_bpe_models:361]\n",
    "    Replacing old number of classes (256) with new number of classes - 256\n",
    "[NeMo I 2022-09-12 23:20:03 ctc_bpe_models:390] Changed tokenizer to ['<unk>', 's', 't', 'e', '▁the', 'd', '▁', '▁a', 'i', 'n', 'a', 'm', '▁to', 'y', 'o', 'ing', '▁and', 'er', 'p', 'u', '▁i\n",
    "n', '▁of', \"'\", '▁i', '▁that', 'ed', 're', 'r', 'c', 'h', 'al', 'ar', 'f', '▁you', '▁s', '▁f', 'an', 'b', '▁it', 'l', 'w', '▁is', '▁p', 'in', '▁we', '▁re', '▁be', 'es', 'g', 'or', '▁he', '▁\n",
    "c', 'ly', 'le', 'k', 'en', '▁for', '▁w', 'll', 'ur', 'ic', 'ri', '▁e', '▁so', 'on', 'ct', 've', '▁b', '▁g', '▁st', 'it', '▁t', '▁do', 'ra', '▁on', '▁was', '▁this', 'ent', 'th', 'ro', 'ce',\n",
    "'▁have', '▁de', '▁o', 'ter', '▁ma', '▁se', '▁co', '▁di', 'ation', '▁with', '▁not', '▁m', 'il', '▁me', 'us', 'ir', '▁are', 'v', '▁but', '▁pro', '▁th', 'ch', '▁con', 'ate', 'me', 'at', 'la',\n",
    "'li', '▁they', 'ver', '▁go', '▁what', '▁ha', 'vi', '▁ne', '▁or', 'ive', '▁as', '▁there', '▁know', 'ment', 'un', 'lo', '▁su', '▁can', 'is', '▁ex', '▁ch', '▁mo', 'ck', 'ul', '▁like', 'tion',\n",
    "'el', '▁le', '▁one', 'ng', 'ci', '▁ca', '▁an', '▁all', 'ne', 'ge', '▁lo', 'x', 'ut', '▁la', '▁if', '▁at', '▁un', 'ol', 'qu', '▁no', '▁fa', 'as', '▁ho', 'ity', '▁just', '▁would', '▁about', '\n",
    "▁from', '▁ba', '▁v', 'mp', '▁think', '▁my', 'z', 'co', 'ad', '▁us', '▁will', '▁li', 'end', '▁by', 'ight', '▁some', '▁po', '▁his', 'ig', 'ry', '▁your', '▁our', '▁out', '▁pa', 'ff', '▁don', '\n",
    "ru', '▁had', '▁te', '▁up', 'j', '▁when', '▁because', '▁which', '▁da', '▁get', 'age', '▁sp', '▁two', '▁bo', '▁say', 'sion', 'ction', '▁pre', '▁were', 'ence', '▁how', '▁time', '▁k', '▁who', '\n",
    "▁mi', '▁right', '▁comp', 'able', '▁she', '▁any', '▁more', 'ugh', '▁now', '▁other', '▁yeah', '▁app', 'ance', '▁uh', '▁also', '▁people', '▁part', '▁want', '▁very', 'ound', '▁work', '▁look', '\n",
    "▁comm', 'port', '▁year', '▁case', '▁court', '▁really', '▁said', 'side', '▁where', '▁could', '▁make', '▁even', '▁dr', '▁every', '▁those', '▁take', '▁ju', '▁three', '▁good', '▁first', '▁shoul\n",
    "d', '▁point', 'q'] vocabulary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "897b8f59",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "##s\r\n",
      "##t\r\n",
      "##e\r\n"
     ]
    }
   ],
   "source": [
    "!head -3 /mnt/data-ssd-1/tmp/tao/results/conformer/train/checkpoints/conformer-tokenizer/vocab.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9c809987",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\"audio_filepath\": \"/data/an4_converted/wavs/an251-fash-b.wav\", \"duration\": 1.0, \"text\": \"yes\"}\r\n",
      "{\"audio_filepath\": \"/data/an4_converted/wavs/an253-fash-b.wav\", \"duration\": 0.7, \"text\": \"go\"}\r\n",
      "{\"audio_filepath\": \"/data/an4_converted/wavs/an254-fash-b.wav\", \"duration\": 0.9, \"text\": \"yes\"}\r\n"
     ]
    }
   ],
   "source": [
    "!head -3 /mnt/data-ssd-1/tmp/tao/data/an4_converted/train_manifest.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bb2802f9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.\r\n",
      "# TLT spec file for fine-tuning a previously trained ASR models based on CTC over the MCV Russian dataset.\r\n",
      "\r\n",
      "trainer:\r\n",
      "  max_epochs: 3   # This is low for demo purposes\r\n",
      "\r\n",
      "tlt_checkpoint_interval: 1\r\n",
      "\r\n",
      "# Whether or not to change the decoder vocabulary.\r\n",
      "# Note that this MUST be set if the labels change, e.g. to a different language's character set\r\n",
      "# or if additional punctuation characters are added.\r\n",
      "change_vocabulary: true\r\n",
      "\r\n",
      "tokenizer:\r\n",
      "  dir: ???\r\n",
      "  type: \"bpe\"  # Can be either bpe or wpe\r\n",
      "\r\n",
      "# Fine-tuning settings: training dataset\r\n",
      "finetuning_ds:\r\n",
      "  manifest_filepath: ???\r\n",
      "  sample_rate: 16000\r\n",
      "  labels: ['<unk>', 's', 't', 'e', '▁the', 'd', '▁', '▁a', 'i', 'n', 'a', 'm', '▁to', 'y', 'o', 'ing', '▁and', 'er', 'p', 'u', '▁i\r\n",
      "  n', '▁of', \"'\", '▁i', '▁that', 'ed', 're', 'r', 'c', 'h', 'al', 'ar', 'f', '▁you', '▁s', '▁f', 'an', 'b', '▁it', 'l', 'w', '▁is', '▁p', 'in', '▁we', '▁re', '▁be', 'es', 'g', 'or', '▁he', '▁\r\n",
      "  c', 'ly', 'le', 'k', 'en', '▁for', '▁w', 'll', 'ur', 'ic', 'ri', '▁e', '▁so', 'on', 'ct', 've', '▁b', '▁g', '▁st', 'it', '▁t', '▁do', 'ra', '▁on', '▁was', '▁this', 'ent', 'th', 'ro', 'ce',\r\n",
      "  '▁have', '▁de', '▁o', 'ter', '▁ma', '▁se', '▁co', '▁di', 'ation', '▁with', '▁not', '▁m', 'il', '▁me', 'us', 'ir', '▁are', 'v', '▁but', '▁pro', '▁th', 'ch', '▁con', 'ate', 'me', 'at', 'la',\r\n",
      "  'li', '▁they', 'ver', '▁go', '▁what', '▁ha', 'vi', '▁ne', '▁or', 'ive', '▁as', '▁there', '▁know', 'ment', 'un', 'lo', '▁su', '▁can', 'is', '▁ex', '▁ch', '▁mo', 'ck', 'ul', '▁like', 'tion',\r\n",
      "  'el', '▁le', '▁one', 'ng', 'ci', '▁ca', '▁an', '▁all', 'ne', 'ge', '▁lo', 'x', 'ut', '▁la', '▁if', '▁at', '▁un', 'ol', 'qu', '▁no', '▁fa', 'as', '▁ho', 'ity', '▁just', '▁would', '▁about', '\r\n",
      "  ▁from', '▁ba', '▁v', 'mp', '▁think', '▁my', 'z', 'co', 'ad', '▁us', '▁will', '▁li', 'end', '▁by', 'ight', '▁some', '▁po', '▁his', 'ig', 'ry', '▁your', '▁our', '▁out', '▁pa', 'ff', '▁don', '\r\n",
      "  ru', '▁had', '▁te', '▁up', 'j', '▁when', '▁because', '▁which', '▁da', '▁get', 'age', '▁sp', '▁two', '▁bo', '▁say', 'sion', 'ction', '▁pre', '▁were', 'ence', '▁how', '▁time', '▁k', '▁who', '\r\n",
      "  ▁mi', '▁right', '▁comp', 'able', '▁she', '▁any', '▁more', 'ugh', '▁now', '▁other', '▁yeah', '▁app', 'ance', '▁uh', '▁also', '▁people', '▁part', '▁want', '▁very', 'ound', '▁work', '▁look', '\r\n",
      "  ▁comm', 'port', '▁year', '▁case', '▁court', '▁really', '▁said', 'side', '▁where', '▁could', '▁make', '▁even', '▁dr', '▁every', '▁those', '▁take', '▁ju', '▁three', '▁good', '▁first', '▁shoul\r\n",
      "  d', '▁point', 'q']\r\n",
      "  \r\n",
      "  batch_size: 32\r\n",
      "  trim_silence: true\r\n",
      "  max_duration: 16.7\r\n",
      "  shuffle: true\r\n",
      "  is_tarred: false\r\n",
      "  tarred_audio_filepaths: null\r\n",
      "\r\n",
      "# Fine-tuning settings: validation dataset\r\n",
      "validation_ds:\r\n",
      "  manifest_filepath: ???\r\n",
      "  sample_rate: 16000\r\n",
      "  labels: ['<unk>', 's', 't', 'e', '▁the', 'd', '▁', '▁a', 'i', 'n', 'a', 'm', '▁to', 'y', 'o', 'ing', '▁and', 'er', 'p', 'u', '▁i\r\n",
      "  n', '▁of', \"'\", '▁i', '▁that', 'ed', 're', 'r', 'c', 'h', 'al', 'ar', 'f', '▁you', '▁s', '▁f', 'an', 'b', '▁it', 'l', 'w', '▁is', '▁p', 'in', '▁we', '▁re', '▁be', 'es', 'g', 'or', '▁he', '▁\r\n",
      "  c', 'ly', 'le', 'k', 'en', '▁for', '▁w', 'll', 'ur', 'ic', 'ri', '▁e', '▁so', 'on', 'ct', 've', '▁b', '▁g', '▁st', 'it', '▁t', '▁do', 'ra', '▁on', '▁was', '▁this', 'ent', 'th', 'ro', 'ce',\r\n",
      "  '▁have', '▁de', '▁o', 'ter', '▁ma', '▁se', '▁co', '▁di', 'ation', '▁with', '▁not', '▁m', 'il', '▁me', 'us', 'ir', '▁are', 'v', '▁but', '▁pro', '▁th', 'ch', '▁con', 'ate', 'me', 'at', 'la',\r\n",
      "  'li', '▁they', 'ver', '▁go', '▁what', '▁ha', 'vi', '▁ne', '▁or', 'ive', '▁as', '▁there', '▁know', 'ment', 'un', 'lo', '▁su', '▁can', 'is', '▁ex', '▁ch', '▁mo', 'ck', 'ul', '▁like', 'tion',\r\n",
      "  'el', '▁le', '▁one', 'ng', 'ci', '▁ca', '▁an', '▁all', 'ne', 'ge', '▁lo', 'x', 'ut', '▁la', '▁if', '▁at', '▁un', 'ol', 'qu', '▁no', '▁fa', 'as', '▁ho', 'ity', '▁just', '▁would', '▁about', '\r\n",
      "  ▁from', '▁ba', '▁v', 'mp', '▁think', '▁my', 'z', 'co', 'ad', '▁us', '▁will', '▁li', 'end', '▁by', 'ight', '▁some', '▁po', '▁his', 'ig', 'ry', '▁your', '▁our', '▁out', '▁pa', 'ff', '▁don', '\r\n",
      "  ru', '▁had', '▁te', '▁up', 'j', '▁when', '▁because', '▁which', '▁da', '▁get', 'age', '▁sp', '▁two', '▁bo', '▁say', 'sion', 'ction', '▁pre', '▁were', 'ence', '▁how', '▁time', '▁k', '▁who', '\r\n",
      "  ▁mi', '▁right', '▁comp', 'able', '▁she', '▁any', '▁more', 'ugh', '▁now', '▁other', '▁yeah', '▁app', 'ance', '▁uh', '▁also', '▁people', '▁part', '▁want', '▁very', 'ound', '▁work', '▁look', '\r\n",
      "  ▁comm', 'port', '▁year', '▁case', '▁court', '▁really', '▁said', 'side', '▁where', '▁could', '▁make', '▁even', '▁dr', '▁every', '▁those', '▁take', '▁ju', '▁three', '▁good', '▁first', '▁shoul\r\n",
      "  d', '▁point', 'q']\r\n",
      "  \r\n",
      "  batch_size: 32\r\n",
      "  shuffle: false\r\n",
      "\r\n",
      "# Fine-tuning settings: optimizer\r\n",
      "optim:\r\n",
      "  name: novograd\r\n",
      "  lr: 0.00000\r\n"
     ]
    }
   ],
   "source": [
    "!cat /mnt/data-ssd-1/tmp/tao/specs/speech_to_text_conformer/finetune_noop.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a6a3fb39",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sample_rate: 16000\r\n",
      "log_prediction: true\r\n",
      "ctc_reduction: mean_batch\r\n",
      "train_ds:\r\n",
      "  manifest_filepath: /data/dataset/tarred_audio_manifest.json\r\n",
      "  sample_rate: 16000\r\n",
      "  batch_size: 16\r\n",
      "  shuffle: true\r\n",
      "  is_tarred: true\r\n",
      "  tarred_audio_filepaths: /data/dataset/audio__OP_0..4095_CL_.tar\r\n",
      "  use_start_end_token: false\r\n",
      "  trim_silence: false\r\n",
      "  max_duration: 20.0\r\n",
      "  min_duration: 0.1\r\n",
      "  num_workers: 16\r\n",
      "  pin_memory: true\r\n",
      "validation_ds:\r\n",
      "  manifest_filepath:\r\n",
      "  - /data/librispeech/LibriSpeech/librivox-test-other.json\r\n",
      "  - /data/librispeech/LibriSpeech/librivox-dev-other.json\r\n",
      "  sample_rate: 16000\r\n",
      "  batch_size: 8\r\n",
      "  shuffle: false\r\n",
      "  use_start_end_token: false\r\n",
      "  num_workers: 8\r\n",
      "test_ds:\r\n",
      "  manifest_filepath: null\r\n",
      "  sample_rate: 16000\r\n",
      "  batch_size: 16\r\n",
      "  shuffle: false\r\n",
      "  num_workers: 8\r\n",
      "  pin_memory: true\r\n",
      "  use_start_end_token: false\r\n",
      "tokenizer:\r\n",
      "  dir: /tokenizers/tokenizer_spe_unigram_v256\r\n",
      "  type: bpe\r\n",
      "  model_path: 8b8f095152034e98b24ab33726708bd0_tokenizer.model\r\n",
      "  vocab_path: 8a5b33d0d7d24579b43eb8b6be27710a_vocab.txt\r\n",
      "  spe_tokenizer_vocab: 98dd8139ae8244d6bfa37d1159bd343a_tokenizer.vocab\r\n",
      "preprocessor:\r\n",
      "  _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor\r\n",
      "  sample_rate: 16000\r\n",
      "  normalize: per_feature\r\n",
      "  window_size: 0.025\r\n",
      "  window_stride: 0.01\r\n",
      "  window: hann\r\n",
      "  features: 80\r\n",
      "  n_fft: 512\r\n",
      "  log: true\r\n",
      "  frame_splicing: 1\r\n",
      "  dither: 1.0e-05\r\n",
      "  pad_to: 0\r\n",
      "  pad_value: 0.0\r\n",
      "spec_augment:\r\n",
      "  _target_: nemo.collections.asr.modules.SpectrogramAugmentation\r\n",
      "  freq_masks: 2\r\n",
      "  time_masks: 10\r\n",
      "  freq_width: 27\r\n",
      "  time_width: 0.05\r\n",
      "encoder:\r\n",
      "  _target_: nemo.collections.asr.modules.ConformerEncoder\r\n",
      "  feat_in: 80\r\n",
      "  feat_out: -1\r\n",
      "  n_layers: 18\r\n",
      "  d_model: 512\r\n",
      "  subsampling: striding\r\n",
      "  subsampling_factor: 4\r\n",
      "  subsampling_conv_channels: -1\r\n",
      "  ff_expansion_factor: 4\r\n",
      "  self_attention_model: rel_pos\r\n",
      "  n_heads: 8\r\n",
      "  att_context_size:\r\n",
      "  - -1\r\n",
      "  - -1\r\n",
      "  xscaling: true\r\n",
      "  untie_biases: true\r\n",
      "  pos_emb_max_len: 5000\r\n",
      "  conv_kernel_size: 31\r\n",
      "  dropout: 0.1\r\n",
      "  dropout_emb: 0.0\r\n",
      "  dropout_att: 0.1\r\n",
      "decoder:\r\n",
      "  _target_: nemo.collections.asr.modules.ConvASRDecoder\r\n",
      "  feat_in: 512\r\n",
      "  num_classes: 256\r\n",
      "  vocabulary:\r\n",
      "  - <unk>\r\n",
      "  - s\r\n",
      "  - t\r\n",
      "  - e\r\n",
      "  - ▁the\r\n",
      "  - d\r\n",
      "  - ▁\r\n",
      "  - ▁a\r\n",
      "  - i\r\n",
      "  - 'n'\r\n",
      "  - a\r\n",
      "  - m\r\n",
      "  - ▁to\r\n",
      "  - 'y'\r\n",
      "  - o\r\n",
      "  - ing\r\n",
      "  - ▁and\r\n",
      "  - er\r\n",
      "  - p\r\n",
      "  - u\r\n",
      "  - ▁in\r\n",
      "  - ▁of\r\n",
      "  - ''''\r\n",
      "  - ▁i\r\n",
      "  - ▁that\r\n",
      "  - ed\r\n",
      "  - re\r\n",
      "  - r\r\n",
      "  - c\r\n",
      "  - h\r\n",
      "  - al\r\n",
      "  - ar\r\n",
      "  - f\r\n",
      "  - ▁you\r\n",
      "  - ▁s\r\n",
      "  - ▁f\r\n",
      "  - an\r\n",
      "  - b\r\n",
      "  - ▁it\r\n",
      "  - l\r\n",
      "  - w\r\n",
      "  - ▁is\r\n",
      "  - ▁p\r\n",
      "  - in\r\n",
      "  - ▁we\r\n",
      "  - ▁re\r\n",
      "  - ▁be\r\n",
      "  - es\r\n",
      "  - g\r\n",
      "  - or\r\n",
      "  - ▁he\r\n",
      "  - ▁c\r\n",
      "  - ly\r\n",
      "  - le\r\n",
      "  - k\r\n",
      "  - en\r\n",
      "  - ▁for\r\n",
      "  - ▁w\r\n",
      "  - ll\r\n",
      "  - ur\r\n",
      "  - ic\r\n",
      "  - ri\r\n",
      "  - ▁e\r\n",
      "  - ▁so\r\n",
      "  - 'on'\r\n",
      "  - ct\r\n",
      "  - ve\r\n",
      "  - ▁b\r\n",
      "  - ▁g\r\n",
      "  - ▁st\r\n",
      "  - it\r\n",
      "  - ▁t\r\n",
      "  - ▁do\r\n",
      "  - ra\r\n",
      "  - ▁on\r\n",
      "  - ▁was\r\n",
      "  - ▁this\r\n",
      "  - ent\r\n",
      "  - th\r\n",
      "  - ro\r\n",
      "  - ce\r\n",
      "  - ▁have\r\n",
      "  - ▁de\r\n",
      "  - ▁o\r\n",
      "  - ter\r\n",
      "  - ▁ma\r\n",
      "  - ▁se\r\n",
      "  - ▁co\r\n",
      "  - ▁di\r\n",
      "  - ation\r\n",
      "  - ▁with\r\n",
      "  - ▁not\r\n",
      "  - ▁m\r\n",
      "  - il\r\n",
      "  - ▁me\r\n",
      "  - us\r\n",
      "  - ir\r\n",
      "  - ▁are\r\n",
      "  - v\r\n",
      "  - ▁but\r\n",
      "  - ▁pro\r\n",
      "  - ▁th\r\n",
      "  - ch\r\n",
      "  - ▁con\r\n",
      "  - ate\r\n",
      "  - me\r\n",
      "  - at\r\n",
      "  - la\r\n",
      "  - li\r\n",
      "  - ▁they\r\n",
      "  - ver\r\n",
      "  - ▁go\r\n",
      "  - ▁what\r\n",
      "  - ▁ha\r\n",
      "  - vi\r\n",
      "  - ▁ne\r\n",
      "  - ▁or\r\n",
      "  - ive\r\n",
      "  - ▁as\r\n",
      "  - ▁there\r\n",
      "  - ▁know\r\n",
      "  - ment\r\n",
      "  - un\r\n",
      "  - lo\r\n",
      "  - ▁su\r\n",
      "  - ▁can\r\n",
      "  - is\r\n",
      "  - ▁ex\r\n",
      "  - ▁ch\r\n",
      "  - ▁mo\r\n",
      "  - ck\r\n",
      "  - ul\r\n",
      "  - ▁like\r\n",
      "  - tion\r\n",
      "  - el\r\n",
      "  - ▁le\r\n",
      "  - ▁one\r\n",
      "  - ng\r\n",
      "  - ci\r\n",
      "  - ▁ca\r\n",
      "  - ▁an\r\n",
      "  - ▁all\r\n",
      "  - ne\r\n",
      "  - ge\r\n",
      "  - ▁lo\r\n",
      "  - x\r\n",
      "  - ut\r\n",
      "  - ▁la\r\n",
      "  - ▁if\r\n",
      "  - ▁at\r\n",
      "  - ▁un\r\n",
      "  - ol\r\n",
      "  - qu\r\n",
      "  - ▁no\r\n",
      "  - ▁fa\r\n",
      "  - as\r\n",
      "  - ▁ho\r\n",
      "  - ity\r\n",
      "  - ▁just\r\n",
      "  - ▁would\r\n",
      "  - ▁about\r\n",
      "  - ▁from\r\n",
      "  - ▁ba\r\n",
      "  - ▁v\r\n",
      "  - mp\r\n",
      "  - ▁think\r\n",
      "  - ▁my\r\n",
      "  - z\r\n",
      "  - co\r\n",
      "  - ad\r\n",
      "  - ▁us\r\n",
      "  - ▁will\r\n",
      "  - ▁li\r\n",
      "  - end\r\n",
      "  - ▁by\r\n",
      "  - ight\r\n",
      "  - ▁some\r\n",
      "  - ▁po\r\n",
      "  - ▁his\r\n",
      "  - ig\r\n",
      "  - ry\r\n",
      "  - ▁your\r\n",
      "  - ▁our\r\n",
      "  - ▁out\r\n",
      "  - ▁pa\r\n",
      "  - ff\r\n",
      "  - ▁don\r\n",
      "  - ru\r\n",
      "  - ▁had\r\n",
      "  - ▁te\r\n",
      "  - ▁up\r\n",
      "  - j\r\n",
      "  - ▁when\r\n",
      "  - ▁because\r\n",
      "  - ▁which\r\n",
      "  - ▁da\r\n",
      "  - ▁get\r\n",
      "  - age\r\n",
      "  - ▁sp\r\n",
      "  - ▁two\r\n",
      "  - ▁bo\r\n",
      "  - ▁say\r\n",
      "  - sion\r\n",
      "  - ction\r\n",
      "  - ▁pre\r\n",
      "  - ▁were\r\n",
      "  - ence\r\n",
      "  - ▁how\r\n",
      "  - ▁time\r\n",
      "  - ▁k\r\n",
      "  - ▁who\r\n",
      "  - ▁mi\r\n",
      "  - ▁right\r\n",
      "  - ▁comp\r\n",
      "  - able\r\n",
      "  - ▁she\r\n",
      "  - ▁any\r\n",
      "  - ▁more\r\n",
      "  - ugh\r\n",
      "  - ▁now\r\n",
      "  - ▁other\r\n",
      "  - ▁yeah\r\n",
      "  - ▁app\r\n",
      "  - ance\r\n",
      "  - ▁uh\r\n",
      "  - ▁also\r\n",
      "  - ▁people\r\n",
      "  - ▁part\r\n",
      "  - ▁want\r\n",
      "  - ▁very\r\n",
      "  - ound\r\n",
      "  - ▁work\r\n",
      "  - ▁look\r\n",
      "  - ▁comm\r\n",
      "  - port\r\n",
      "  - ▁year\r\n",
      "  - ▁case\r\n",
      "  - ▁court\r\n",
      "  - ▁really\r\n",
      "  - ▁said\r\n",
      "  - side\r\n",
      "  - ▁where\r\n",
      "  - ▁could\r\n",
      "  - ▁make\r\n",
      "  - ▁even\r\n",
      "  - ▁dr\r\n",
      "  - ▁every\r\n",
      "  - ▁those\r\n",
      "  - ▁take\r\n",
      "  - ▁ju\r\n",
      "  - ▁three\r\n",
      "  - ▁good\r\n",
      "  - ▁first\r\n",
      "  - ▁should\r\n",
      "  - ▁point\r\n",
      "  - q\r\n",
      "optim:\r\n",
      "  name: adamw\r\n",
      "  lr: 2.0\r\n",
      "  betas:\r\n",
      "  - 0.9\r\n",
      "  - 0.98\r\n",
      "  weight_decay: 0.001\r\n",
      "  sched:\r\n",
      "    name: NoamAnnealing\r\n",
      "    d_model: 512\r\n",
      "    warmup_steps: 10000\r\n",
      "    warmup_ratio: null\r\n",
      "    min_lr: 1.0e-06\r\n",
      "target: nemo.collections.asr.models.ctc_bpe_models.EncDecCTCModelBPE\r\n",
      "nemo_version: 1.9.0\r\n"
     ]
    }
   ],
   "source": [
    "!cat /home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/model_config.yaml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6055f9a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "docker cp /home/georg/models/speechtotext_en_us_conformer_vtrainable_v4.0/encrypted_model_weights.ckpt 2c25801c:/data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55fee8d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "md = \"/results/conformer/finetune/checkpoints/\"\n",
    "d = torch.load([md + fn for fn in os.listdir(md) if \"ckpt\" in fn][0], map_location=\"cpu\")\n",
    "d2 = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")\n",
    "torch.save(d2[\"state_dict\"], \"/results/conformer/finetune/checkpoints/decrypted.pt\")\n",
    "torch.load(\"/data/encrypted_model_weights.ckpt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2869af72",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/data/encrypted_model_weights.ckpt\", \"rb\") as f:\n",
    "    b = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4becdf42",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from core import checkpoint_encryption\n",
    "\n",
    "md = \"/results/conformer/finetune/checkpoints/\"\n",
    "d = torch.load([md + fn for fn in os.listdir(md) if \"ckpt\" in fn][0], map_location=\"cpu\")\n",
    "\n",
    "with open(\"/data/encrypted_model_weights.ckpt\", \"rb\") as f:\n",
    "    b = f.read()\n",
    "d[\"state_dict\"] = b\n",
    "d[\"state_dict_encrypted\"] = \"qahxrkjnrz\"\n",
    "# d = {\"state_dict\": b, \"state_dict_encrypted\": \"qahxrkjnrz\"}\n",
    "d2 = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a034520",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"/data/encrypted_model_weights.ckpt\", \"rb\") as f:\n",
    "with gzip.open('/data/encrypted_model_weights.ckpt', 'rb') as f:\n",
    "    b = f.read()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa8ddec6",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e4ba54",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1add0708",
   "metadata": {},
   "outputs": [],
   "source": [
    "d2 = checkpoint_encryption.decrypt_checkpoint(d, \"tlt_encode\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a069fb15",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet evaluate \\\n",
    "     -e $SPECS_DIR/speech_to_text_conformer/evaluate.yaml \\\n",
    "     -g 1 \\\n",
    "     -k $KEY \\\n",
    "     -m $RESULTS_DIR/conformer/train/checkpoints/speechtotext_en_us_conformer.tlt \\\n",
    "     -r $RESULTS_DIR/conformer/evaluate \\\n",
    "     test_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30df3105",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18419882",
   "metadata": {},
   "outputs": [],
   "source": [
    "tao speech_to_text_citrinet evaluate \\\n",
    "     -e $SPECS_DIR/speech_to_text_conformer/evaluate.yaml \\\n",
    "     -g 1 \\\n",
    "     -k $KEY \\\n",
    "     -m $RESULTS_DIR/tmp/finetuned-model_epoch_11.tlt \\\n",
    "     -r $RESULTS_DIR/conformer/evaluate \\\n",
    "     test_ds.manifest_filepath=$DATA_DIR/an4_converted/test_manifest.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7dec0d7",
   "metadata": {},
   "outputs": [],
   "source": [
    "/home/jenkins/agent/workspace/tlt-pytorch-main-nightly/core/checkpoint_encryption.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01c37537",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33e9801d",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46b2b198",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "412e7a41",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5b7c258",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b216d80",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a8955f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}