diff --git a/batch_demo.ipynb b/batch_demo.ipynb
index ef00ac9..58e9de0 100644
--- a/batch_demo.ipynb
+++ b/batch_demo.ipynb
@@ -1 +1,1098 @@
-{"cells":[{"cell_type":"markdown","id":"c414b951-1124-47b9-8413-bf86e83037fc","metadata":{},"source":"## Instructions to setup a demo run"},{"cell_type":"markdown","id":"4ced3393-884f-4cfd-a41e-10847474a044","metadata":{},"source":"1. Prepare an input file similar to csv files given in ./data/demo/\n2. Define the path to the input file in below cell\n3. Define the parameter for prediction\n4. Run the cell to get output file\n5. The last five columns of the output csv file saved will have the prediction results"},{"cell_type":"code","execution_count":3,"id":"3c1841bb-8c35-4562-b86f-30777bb299bb","metadata":{},"outputs":[],"source":"parameter = 'kcat' # allowed values: [\"kcat\", \"Km\", \"Ki\"] \nparameter = parameter.lower()\n\nuse_cpu = 1 # set to 0 if you have GPU enabled\n\ninput_file_path = './demo/batch_kcat.csv'"},{"cell_type":"markdown","id":"db40ad26-0481-405b-8323-70d23f0ef44e","metadata":{},"source":"## Navigate to below cell and click \"Run->Run Selected Cell\" to get prediction"},{"cell_type":"markdown","id":"81de60cf-39b5-4546-83da-497c15e01b59","metadata":{},"source":"The result will be printed on the right column"},{"cell_type":"code","execution_count":null,"id":"d31b0d48-aa79-4701-a535-9104a83885a2","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":3,"id":"fcf67818-1d36-48a1-bac1-ee2d125610b4","metadata":{},"outputs":[],"source":"parameter = 'kcat' # allowed values: [\"kcat\", \"Km\", \"Ki\"] \nparameter = parameter.lower()\n\nuse_cpu = 1 # set to 0 if you have GPU enabled\n\ninput_file_path = './demo/batch_kcat.csv'"},{"cell_type":"code","execution_count":null,"id":"13857d8e-efa2-42ac-9f66-1c27ea49a07f","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":null,"id":"d89fddab-dd64-452e-bb86-e0e0aa8cd214","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":null,"id":"ee02f79f-20c3-4e63-ac02-f8cf47026c1d","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":7,"id":"d737ae1b-3a84-4995-8cda-9d2905f2b06d","metadata":{"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[{"name":"stdout","output_type":"stream","text":["Predicting.. This will take a while..\n","\n"]},{"name":"stderr","output_type":"stream","text":["gzip: ./demo/batch_kcat_input.json.gz already exists;\tnot overwritten\n","/opt/conda/lib/python3.12/site-packages/rotary_embedding_torch/rotary_embedding_torch.py:35: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n","  @autocast(enabled = False)\n","/opt/conda/lib/python3.12/site-packages/rotary_embedding_torch/rotary_embedding_torch.py:268: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n","  @autocast(enabled = False)\n","/root/capsule/code/catpred/utils.py:501: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  vars(torch.load(path, map_location=lambda storage, loc: storage)[\"args\"]),\n","0it [00:00, ?it/s]"]},{"name":"stdout","output_type":"stream","text":["calculating protein embed only on cpu\n","Loading training args\n","Loading models\n","Setting molecule featurization parameters to default.\n","Loading data\n"]},{"name":"stderr","output_type":"stream","text":["9it [01:31,  8.74s/it]/root/capsule/code/catpred/data/cache_utils.py:95: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  return torch.load(str(entry_path))\n","14it [01:31,  6.55s/it]\n","100%|██████████| 14/14 [00:00<00:00, 58428.12it/s]\n","100%|██████████| 14/14 [00:00<00:00, 1643.49it/s]\n","  0%|          | 0/10 [00:00<?, ?it/s]"]},{"name":"stdout","output_type":"stream","text":["Validating SMILES\n","Test size = 14\n"]},{"name":"stderr","output_type":"stream","text":["/root/capsule/code/catpred/utils.py:113: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  state = torch.load(path, map_location=lambda storage, loc: storage)\n","/root/capsule/code/catpred/utils.py:446: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  state = torch.load(path, map_location=lambda storage, loc: storage)\n","\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","100%|██████████| 1/1 [00:01<00:00,  1.61s/it]\u001b[A\n"," 10%|█         | 1/10 [00:04<00:39,  4.39s/it]/root/capsule/code/catpred/utils.py:113: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  state = torch.load(path, map_location=lambda storage, loc: storage)\n","/root/capsule/code/catpred/utils.py:446: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  state = torch.load(path, map_location=lambda storage, loc: storage)\n","\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","100%|██████████| 1/1 [00:00<00:00,  7.27it/s]\u001b[A\n"," 20%|██        | 2/10 [00:04<00:15,  1.97s/it][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.55it/s]\u001b[A\n"," 30%|███       | 3/10 [00:05<00:08,  1.25s/it][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.57it/s]\u001b[A\n"," 40%|████      | 4/10 [00:05<00:05,  1.17it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.53it/s]\u001b[A\n"," 50%|█████     | 5/10 [00:05<00:03,  1.55it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.57it/s]\u001b[A\n"," 60%|██████    | 6/10 [00:05<00:02,  1.95it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.56it/s]\u001b[A\n"," 70%|███████   | 7/10 [00:06<00:01,  2.32it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.66it/s]\u001b[A\n"," 80%|████████  | 8/10 [00:06<00:00,  2.66it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.65it/s]\u001b[A\n"," 90%|█████████ | 9/10 [00:06<00:00,  2.95it/s][A"]},{"name":"stdout","output_type":"stream","text":["Creating protein model\n","MoleculeModel(\n","  (softplus): Softplus(beta=1.0, threshold=20.0)\n","  (encoder): MPN(\n","    (encoder): ModuleList(\n","      (0): MPNEncoder(\n","        (dropout): Dropout(p=0.0, inplace=False)\n","        (act_func): ReLU()\n","        (W_i): Linear(in_features=147, out_features=300, bias=False)\n","        (W_h): Linear(in_features=300, out_features=300, bias=False)\n","        (W_o): Linear(in_features=433, out_features=300, bias=True)\n","      )\n","    )\n","  )\n","  (seq_embedder): Embedding(21, 36, padding_idx=20)\n","  (rotary_embedder): RotaryEmbedding()\n","  (multihead_attn): MultiheadAttention(\n","    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n","  )\n","  (attentive_pooler): AttentivePooling(\n","    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n","    (tanh): Tanh()\n","    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n","    (softmax): Softmax(dim=1)\n","  )\n","  (readout): Sequential(\n","    (0): Dropout(p=0.0, inplace=False)\n","    (1): Linear(in_features=1616, out_features=300, bias=True)\n","    (2): ReLU()\n","    (3): Dropout(p=0.0, inplace=False)\n","    (4): Linear(in_features=300, out_features=2, bias=True)\n","  )\n",")\n","Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n","Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n","Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n","Loading pretrained parameter \"seq_embedder.weight\".\n","Loading pretrained parameter \"rotary_embedder.freqs\".\n","Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n","Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n","Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n","Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n","Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n","Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n","Loading pretrained parameter \"readout.1.weight\".\n","Loading pretrained parameter \"readout.1.bias\".\n","Loading pretrained parameter \"readout.4.weight\".\n","Loading pretrained parameter \"readout.4.bias\".\n","Moving model to cuda\n","Saving predictions to ./demo/batch_kcat_input_output.csv\n","Elapsed time = 0:01:39\n","Prediction completed.\n"]},{"name":"stderr","output_type":"stream","text":["\n","  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n","100%|██████████| 1/1 [00:00<00:00,  7.65it/s]\u001b[A\n","100%|██████████| 10/10 [00:06<00:00,  1.46it/s]A\n"]},{"name":"stdout","output_type":"stream","text":["Output saved to ./demo/batch_kcat_input_output.csv\n"]}],"source":"import os\nimport pandas as pd\nimport numpy as np\nfrom IPython.display import Image, display\nfrom rdkit import Chem\nfrom IPython.display import display, Latex, Math\n\ndef create_csv_sh(parameter, input_file_path):\n    df = pd.read_csv(input_file_path)\n    smiles_list = df.SMILES\n    seq_list = df.sequence\n    smiles_list_new = []\n    i=0\n    for smi in smiles_list:\n        try:\n            mol = Chem.MolFromSmiles(smi)\n            smi = Chem.MolToSmiles(mol)\n        except:\n            print(f'Invalid SMILES input in input row {i}')\n            print('Correct your input! Exiting..')\n            return\n        if parameter=='kcat':\n            if '.' in smi:\n              x = smi.split('.')\n              y = sorted(x)\n              smi = '.'.join(y)\n        smiles_list_new.append(smi)\n        i+=1\n  \n    i=0\n    valid_aas = list('ACDEFGHIKLMNPQRSTVWY')\n    for seq in seq_list:\n      for aa in seq:\n        if not aa in valid_aas:\n          print(f'Invalid Enzyme sequence input in row {i}!')\n          print('Correct your input! Exiting..')\n          return\n      i+=1\n\n    input_file_new_path = f'{input_file_path[:-4]}_input.csv'\n    df['SMILES'] = smiles_list_new\n    df.to_csv(input_file_new_path)\n    \n    f = open(f'predict.sh', 'w')\n    f.write(f'''\n    TEST_FILE_PREFIX={input_file_new_path[:-4]}\n    RECORDS_FILE=${{TEST_FILE_PREFIX}}.json\n    CHECKPOINT_DIR=../data/pretrained/production/{parameter}/\n    \n    python ./scripts/create_pdbrecords.py --data_file ${{TEST_FILE_PREFIX}}.csv --out_file ${{RECORDS_FILE}}\n    gzip ${{RECORDS_FILE}}\n    python predict.py --test_path ${{TEST_FILE_PREFIX}}.csv --preds_path ${{TEST_FILE_PREFIX}}_output.csv --checkpoint_dir $CHECKPOINT_DIR --uncertainty_method mve --smiles_column SMILES --individual_ensemble_predictions --protein_records_path ${{RECORDS_FILE}}.gz\n    ''')\n    f.close()\n    \n    return input_file_new_path[:-4]+'_output.csv'\n\noutfile = create_csv_sh(parameter, input_file_path)\n\nprint('Predicting.. This will take a while..\\n')\n\nif use_cpu:\n    os.system(\"export PROTEIN_EMBED_USE_CPU=1;./predict.sh\")\nelse:\n    os.system(\"export PROTEIN_EMBED_USE_CPU=0;./predict.sh\") #>/dev/null 2>&1\n\ndef get_predictions(parameter, outfile):\n    df = pd.read_csv(outfile)\n    pred_col = []\n    pred_logcol = []\n    pred_sd_totcol = []\n    pred_sd_aleacol = []\n    pred_sd_epicol = []\n    \n    for ind, row in df.iterrows():\n        unit = 'mM'\n        if parameter=='kcat':\n            parameter_print = 'k_{cat}'\n            parameter_print_log = 'log_{10}(k_{cat})'\n            target_col = 'log10kcat_max'\n            unit = 's^(-1)'\n        elif parameter=='km':\n            target_col = 'log10km_mean'\n            parameter_print = 'K_{m}'\n            parameter_print_log = 'log_{10}(K_{m})'\n        else:\n            target_col = 'log10ki_mean'\n            parameter_print = 'K_{i}'\n            parameter_print_log = 'log_{10}(K_{i})'\n    \n        unc_col = f'{target_col}_mve_uncal_var'\n        model_cols = [col for col in df.columns if col.startswith(target_col) and 'model_' in col]\n    \n        unc = df[unc_col].iloc[0]\n    \n        prediction = df[target_col].iloc[0]\n        prediction_linear = np.power(10, prediction)\n    \n        model_out = df[target_col].iloc[0]\n        model_outs = np.array([df[col].iloc[0] for col in model_cols])\n        # print(model_outs)\n        epi_unc = np.var(model_outs)#np.sum(np.power(2, model_outs))/10. - np.power(2, model_out)\n        alea_unc = unc - epi_unc\n        epi_unc = np.sqrt(epi_unc)\n        alea_unc = np.sqrt(alea_unc)\n        unc = np.sqrt(unc)\n        \n        # print(unc-epi_unc-alea_unc)\n        # def display_outs(prediction_type, out, alea_output, epi_output, unit):\n        pred_col.append(prediction_linear)\n        pred_logcol.append(prediction)\n        pred_sd_totcol.append(unc)\n        pred_sd_aleacol.append(alea_unc)\n        pred_sd_epicol.append(epi_unc)\n\n    df[f'Prediction_({unit})'] = pred_col\n    df['Prediction_log10'] = pred_logcol\n    df['SD_total'] = pred_sd_totcol\n    df['SD_aleatoric'] = pred_sd_aleacol\n    df['SD_epistemic'] = pred_sd_epicol\n\n    return df\n\noutput_final = get_predictions(parameter, outfile)\noutput_final.to_csv(f'{outfile}')\nprint('Output saved to', outfile)"},{"cell_type":"code","execution_count":19,"id":"3870129e-5fd5-4eae-b8b7-88bc727a8cf9","metadata":{},"outputs":[{"name":"stdout","output_type":"stream","text":["/opt/conda/lib/python3.12/site-packages/progres/progres.py:411: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  loaded_model = torch.load(trained_model_fp, map_location=device)\n"]}],"source":""},{"cell_type":"code","execution_count":20,"id":"89a27021-0de1-421d-b780-59cf63d7b95a","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":22,"id":"3a395ea2-e968-4e09-a96d-21bad113d0ea","metadata":{},"outputs":[{"name":"stderr","output_type":"stream","text":["/tmp/ipykernel_1562/3523922536.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n","  embeds = torch.load('./demo/structures_embeds.pt')\n"]}],"source":""},{"cell_type":"code","execution_count":25,"id":"102914ef-79c6-4ed6-af20-b121636f2c2b","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":27,"id":"1f35cc65-e070-445c-b369-119ae5085673","metadata":{},"outputs":[],"source":""},{"cell_type":"code","execution_count":null,"id":"c256c0e7-3913-44eb-9786-f08ba77ff7c1","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.12.4"}},"nbformat":4,"nbformat_minor":5}
\ No newline at end of file
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "c414b951-1124-47b9-8413-bf86e83037fc",
+   "metadata": {},
+   "source": [
+    "## Instructions to setup a demo run"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "4ced3393-884f-4cfd-a41e-10847474a044",
+   "metadata": {},
+   "source": [
+    "1. Prepare an input file similar to csv files given in ./data/demo/\n",
+    "2. Define the path to the input file in below cell\n",
+    "3. Define the parameter for prediction\n",
+    "4. Run the cell to get output file\n",
+    "5. The last five columns of the output csv file saved will have the prediction results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "3c1841bb-8c35-4562-b86f-30777bb299bb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "parameter = 'kcat' # allowed values: [\"kcat\", \"Km\", \"Ki\"] \n",
+    "parameter = parameter.lower()\n",
+    "\n",
+    "use_cpu = 1 # set to 0 if you have GPU enabled\n",
+    "\n",
+    "input_file_path = './demo/batch_kcat.csv'"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "db40ad26-0481-405b-8323-70d23f0ef44e",
+   "metadata": {},
+   "source": [
+    "## Navigate to below cell and click \"Run->Run Selected Cell\" to get prediction"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "81de60cf-39b5-4546-83da-497c15e01b59",
+   "metadata": {},
+   "source": [
+    "The result will be printed on the right column"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d31b0d48-aa79-4701-a535-9104a83885a2",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fcf67818-1d36-48a1-bac1-ee2d125610b4",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "parameter = 'kcat' # allowed values: [\"kcat\", \"Km\", \"Ki\"] \n",
+    "parameter = parameter.lower()\n",
+    "\n",
+    "use_cpu = 1 # set to 0 if you have GPU enabled\n",
+    "\n",
+    "input_file_path = './demo/batch_kcat.csv'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13857d8e-efa2-42ac-9f66-1c27ea49a07f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d89fddab-dd64-452e-bb86-e0e0aa8cd214",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ee02f79f-20c3-4e63-ac02-f8cf47026c1d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d737ae1b-3a84-4995-8cda-9d2905f2b06d",
+   "metadata": {
+    "collapsed": true,
+    "jupyter": {
+     "outputs_hidden": true
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Predicting.. This will take a while..\n",
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "gzip: ./demo/batch_kcat_input.json.gz already exists;\tnot overwritten\n",
+      "/opt/conda/lib/python3.12/site-packages/rotary_embedding_torch/rotary_embedding_torch.py:35: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  @autocast(enabled = False)\n",
+      "/opt/conda/lib/python3.12/site-packages/rotary_embedding_torch/rotary_embedding_torch.py:268: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n",
+      "  @autocast(enabled = False)\n",
+      "/root/capsule/code/catpred/utils.py:501: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  vars(torch.load(path, map_location=lambda storage, loc: storage)[\"args\"]),\n",
+      "0it [00:00, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "calculating protein embed only on cpu\n",
+      "Loading training args\n",
+      "Loading models\n",
+      "Setting molecule featurization parameters to default.\n",
+      "Loading data\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "9it [01:31,  8.74s/it]/root/capsule/code/catpred/data/cache_utils.py:95: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  return torch.load(str(entry_path))\n",
+      "14it [01:31,  6.55s/it]\n",
+      "100%|██████████| 14/14 [00:00<00:00, 58428.12it/s]\n",
+      "100%|██████████| 14/14 [00:00<00:00, 1643.49it/s]\n",
+      "  0%|          | 0/10 [00:00<?, ?it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Validating SMILES\n",
+      "Test size = 14\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/root/capsule/code/catpred/utils.py:113: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state = torch.load(path, map_location=lambda storage, loc: storage)\n",
+      "/root/capsule/code/catpred/utils.py:446: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state = torch.load(path, map_location=lambda storage, loc: storage)\n",
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 1/1 [00:01<00:00,  1.61s/it]\u001b[A\n",
+      " 10%|█         | 1/10 [00:04<00:39,  4.39s/it]/root/capsule/code/catpred/utils.py:113: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state = torch.load(path, map_location=lambda storage, loc: storage)\n",
+      "/root/capsule/code/catpred/utils.py:446: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  state = torch.load(path, map_location=lambda storage, loc: storage)\n",
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.27it/s]\u001b[A\n",
+      " 20%|██        | 2/10 [00:04<00:15,  1.97s/it][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.55it/s]\u001b[A\n",
+      " 30%|███       | 3/10 [00:05<00:08,  1.25s/it][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.57it/s]\u001b[A\n",
+      " 40%|████      | 4/10 [00:05<00:05,  1.17it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.53it/s]\u001b[A\n",
+      " 50%|█████     | 5/10 [00:05<00:03,  1.55it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.57it/s]\u001b[A\n",
+      " 60%|██████    | 6/10 [00:05<00:02,  1.95it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.56it/s]\u001b[A\n",
+      " 70%|███████   | 7/10 [00:06<00:01,  2.32it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.66it/s]\u001b[A\n",
+      " 80%|████████  | 8/10 [00:06<00:00,  2.66it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.65it/s]\u001b[A\n",
+      " 90%|█████████ | 9/10 [00:06<00:00,  2.95it/s][A"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Creating protein model\n",
+      "MoleculeModel(\n",
+      "  (softplus): Softplus(beta=1.0, threshold=20.0)\n",
+      "  (encoder): MPN(\n",
+      "    (encoder): ModuleList(\n",
+      "      (0): MPNEncoder(\n",
+      "        (dropout): Dropout(p=0.0, inplace=False)\n",
+      "        (act_func): ReLU()\n",
+      "        (W_i): Linear(in_features=147, out_features=300, bias=False)\n",
+      "        (W_h): Linear(in_features=300, out_features=300, bias=False)\n",
+      "        (W_o): Linear(in_features=433, out_features=300, bias=True)\n",
+      "      )\n",
+      "    )\n",
+      "  )\n",
+      "  (seq_embedder): Embedding(21, 36, padding_idx=20)\n",
+      "  (rotary_embedder): RotaryEmbedding()\n",
+      "  (multihead_attn): MultiheadAttention(\n",
+      "    (out_proj): NonDynamicallyQuantizableLinear(in_features=36, out_features=36, bias=True)\n",
+      "  )\n",
+      "  (attentive_pooler): AttentivePooling(\n",
+      "    (linear1): Linear(in_features=1316, out_features=1316, bias=True)\n",
+      "    (tanh): Tanh()\n",
+      "    (linear2): Linear(in_features=1316, out_features=1, bias=True)\n",
+      "    (softmax): Softmax(dim=1)\n",
+      "  )\n",
+      "  (readout): Sequential(\n",
+      "    (0): Dropout(p=0.0, inplace=False)\n",
+      "    (1): Linear(in_features=1616, out_features=300, bias=True)\n",
+      "    (2): ReLU()\n",
+      "    (3): Dropout(p=0.0, inplace=False)\n",
+      "    (4): Linear(in_features=300, out_features=2, bias=True)\n",
+      "  )\n",
+      ")\n",
+      "Loading pretrained parameter \"encoder.encoder.0.cached_zero_vector\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_i.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_h.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.weight\".\n",
+      "Loading pretrained parameter \"encoder.encoder.0.W_o.bias\".\n",
+      "Loading pretrained parameter \"seq_embedder.weight\".\n",
+      "Loading pretrained parameter \"rotary_embedder.freqs\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.in_proj_bias\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.weight\".\n",
+      "Loading pretrained parameter \"multihead_attn.out_proj.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear1.bias\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.weight\".\n",
+      "Loading pretrained parameter \"attentive_pooler.linear2.bias\".\n",
+      "Loading pretrained parameter \"readout.1.weight\".\n",
+      "Loading pretrained parameter \"readout.1.bias\".\n",
+      "Loading pretrained parameter \"readout.4.weight\".\n",
+      "Loading pretrained parameter \"readout.4.bias\".\n",
+      "Moving model to cuda\n",
+      "Saving predictions to ./demo/batch_kcat_input_output.csv\n",
+      "Elapsed time = 0:01:39\n",
+      "Prediction completed.\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "  0%|          | 0/1 [00:00<?, ?it/s]\u001b[A\n",
+      "100%|██████████| 1/1 [00:00<00:00,  7.65it/s]\u001b[A\n",
+      "100%|██████████| 10/10 [00:06<00:00,  1.46it/s]A\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Output saved to ./demo/batch_kcat_input_output.csv\n"
+     ]
+    }
+   ],
+   "source": [
+    "import os\n",
+    "import pandas as pd\n",
+    "import numpy as np\n",
+    "from IPython.display import Image, display\n",
+    "from rdkit import Chem\n",
+    "from IPython.display import display, Latex, Math\n",
+    "\n",
+    "def create_csv_sh(parameter, input_file_path):\n",
+    "    df = pd.read_csv(input_file_path)\n",
+    "    smiles_list = df.SMILES\n",
+    "    seq_list = df.sequence\n",
+    "    smiles_list_new = []\n",
+    "    i=0\n",
+    "    for smi in smiles_list:\n",
+    "        try:\n",
+    "            mol = Chem.MolFromSmiles(smi)\n",
+    "            smi = Chem.MolToSmiles(mol)\n",
+    "        except:\n",
+    "            print(f'Invalid SMILES input in input row {i}')\n",
+    "            print('Correct your input! Exiting..')\n",
+    "            return\n",
+    "        if parameter=='kcat':\n",
+    "            if '.' in smi:\n",
+    "              x = smi.split('.')\n",
+    "              y = sorted(x)\n",
+    "              smi = '.'.join(y)\n",
+    "        smiles_list_new.append(smi)\n",
+    "        i+=1\n",
+    "  \n",
+    "    i=0\n",
+    "    valid_aas = list('ACDEFGHIKLMNPQRSTVWY')\n",
+    "    for seq in seq_list:\n",
+    "      for aa in seq:\n",
+    "        if not aa in valid_aas:\n",
+    "          print(f'Invalid Enzyme sequence input in row {i}!')\n",
+    "          print('Correct your input! Exiting..')\n",
+    "          return\n",
+    "      i+=1\n",
+    "\n",
+    "    input_file_new_path = f'{input_file_path[:-4]}_input.csv'\n",
+    "    df['SMILES'] = smiles_list_new\n",
+    "    df.to_csv(input_file_new_path)\n",
+    "    \n",
+    "    f = open(f'predict.sh', 'w')\n",
+    "    f.write(f'''\n",
+    "    TEST_FILE_PREFIX={input_file_new_path[:-4]}\n",
+    "    RECORDS_FILE=${{TEST_FILE_PREFIX}}.json\n",
+    "    CHECKPOINT_DIR=../data/pretrained/production/{parameter}/\n",
+    "    \n",
+    "    python ./scripts/create_pdbrecords.py --data_file ${{TEST_FILE_PREFIX}}.csv --out_file ${{RECORDS_FILE}}\n",
+    "    gzip ${{RECORDS_FILE}}\n",
+    "    python predict.py --test_path ${{TEST_FILE_PREFIX}}.csv --preds_path ${{TEST_FILE_PREFIX}}_output.csv --checkpoint_dir $CHECKPOINT_DIR --uncertainty_method mve --smiles_column SMILES --individual_ensemble_predictions --protein_records_path ${{RECORDS_FILE}}.gz\n",
+    "    ''')\n",
+    "    f.close()\n",
+    "    \n",
+    "    return input_file_new_path[:-4]+'_output.csv'\n",
+    "\n",
+    "outfile = create_csv_sh(parameter, input_file_path)\n",
+    "\n",
+    "print('Predicting.. This will take a while..\\n')\n",
+    "\n",
+    "if use_cpu:\n",
+    "    os.system(\"export PROTEIN_EMBED_USE_CPU=1;./predict.sh\")\n",
+    "else:\n",
+    "    os.system(\"export PROTEIN_EMBED_USE_CPU=0;./predict.sh\") #>/dev/null 2>&1\n",
+    "\n",
+    "def get_predictions(parameter, outfile):\n",
+    "    df = pd.read_csv(outfile)\n",
+    "    pred_col = []\n",
+    "    pred_logcol = []\n",
+    "    pred_sd_totcol = []\n",
+    "    pred_sd_aleacol = []\n",
+    "    pred_sd_epicol = []\n",
+    "    \n",
+    "    for ii in  range(df.shape[0]):\n",
+    "        unit = 'mM'\n",
+    "        if parameter=='kcat':\n",
+    "            parameter_print = 'k_{cat}'\n",
+    "            parameter_print_log = 'log_{10}(k_{cat})'\n",
+    "            target_col = 'log10kcat_max'\n",
+    "            unit = 's^(-1)'\n",
+    "        elif parameter=='km':\n",
+    "            target_col = 'log10km_mean'\n",
+    "            parameter_print = 'K_{m}'\n",
+    "            parameter_print_log = 'log_{10}(K_{m})'\n",
+    "        else:\n",
+    "            target_col = 'log10ki_mean'\n",
+    "            parameter_print = 'K_{i}'\n",
+    "            parameter_print_log = 'log_{10}(K_{i})'\n",
+    "    \n",
+    "        unc_col = f'{target_col}_mve_uncal_var'\n",
+    "        model_cols = [col for col in df.columns if col.startswith(target_col) and 'model_' in col]\n",
+    "    \n",
+    "        unc = df[unc_col].iloc[ii]\n",
+    "    \n",
+    "        prediction = df[target_col].iloc[ii]\n",
+    "        prediction_linear = np.power(10, prediction)\n",
+    "    \n",
+    "        model_out = df[target_col].iloc[ii]\n",
+    "        model_outs = np.array([df[col].iloc[ii] for col in model_cols])\n",
+    "        # print(model_outs)\n",
+    "        epi_unc = np.var(model_outs)#np.sum(np.power(2, model_outs))/10. - np.power(2, model_out)\n",
+    "        alea_unc = unc - epi_unc\n",
+    "        epi_unc = np.sqrt(epi_unc)\n",
+    "        alea_unc = np.sqrt(alea_unc)\n",
+    "        unc = np.sqrt(unc)\n",
+    "        \n",
+    "        # print(unc-epi_unc-alea_unc)\n",
+    "        # def display_outs(prediction_type, out, alea_output, epi_output, unit):\n",
+    "        pred_col.append(prediction_linear)\n",
+    "        pred_logcol.append(prediction)\n",
+    "        pred_sd_totcol.append(unc)\n",
+    "        pred_sd_aleacol.append(alea_unc)\n",
+    "        pred_sd_epicol.append(epi_unc)\n",
+    "\n",
+    "    df[f'Prediction_({unit})'] = pred_col\n",
+    "    df['Prediction_log10'] = pred_logcol\n",
+    "    df['SD_total'] = pred_sd_totcol\n",
+    "    df['SD_aleatoric'] = pred_sd_aleacol\n",
+    "    df['SD_epistemic'] = pred_sd_epicol\n",
+    "\n",
+    "    return df\n",
+    "\n",
+    "output_final = get_predictions(parameter, outfile)\n",
+    "output_final.to_csv(f'{outfile}')\n",
+    "print('Output saved to', outfile)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "3870129e-5fd5-4eae-b8b7-88bc727a8cf9",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "/opt/conda/lib/python3.12/site-packages/progres/progres.py:411: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  loaded_model = torch.load(trained_model_fp, map_location=device)\n"
+     ]
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "89a27021-0de1-421d-b780-59cf63d7b95a",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 22,
+   "id": "3a395ea2-e968-4e09-a96d-21bad113d0ea",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/tmp/ipykernel_1562/3523922536.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n",
+      "  embeds = torch.load('./demo/structures_embeds.pt')\n"
+     ]
+    }
+   ],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "102914ef-79c6-4ed6-af20-b121636f2c2b",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 27,
+   "id": "1f35cc65-e070-445c-b369-119ae5085673",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c256c0e7-3913-44eb-9786-f08ba77ff7c1",
+   "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.12.4"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}