fix for multimer representation

sokrypton · Nov 1, 2023 · 29395e7 · 29395e7
1 parent 52cbf8c
commit 29395e7
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/beta/convert_256_to_384_rep.ipynb b/beta/convert_256_to_384_rep.ipynb
@@ -4,7 +4,7 @@
   "metadata": {
     "colab": {
       "provenance": [],
-      "authorship_tag": "ABX9TyNOXTm3jeyevIJtPmlsX2Gb",
+      "authorship_tag": "ABX9TyMxPLX4u9e10BdDErRiKV+Y",
       "include_colab_link": true
     },
     "kernelspec": {
@@ -28,7 +28,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": 14,
+      "execution_count": null,
       "metadata": {
         "id": "SW7PGft_TTOL"
       },
@@ -54,11 +54,15 @@
         "with open(example,\"rb\") as handle:\n",
         "  data = pickle.load(handle)\n",
         "\n",
-        "tmp = example.split(\"_\")\n",
-        "model_num = tmp[-3]\n",
-        "model_type = tmp[-5]\n",
+        "model_num = example.split(\"_\")[-3]\n",
+        "if \"multimer_v3\" in example:\n",
+        "  model_type = \"_multimer_v3\"\n",
+        "elif \"ptm\" in example:\n",
+        "  model_type = \"_ptm\"\n",
+        "else:\n",
+        "  model_type = \"\"\n",
         "print(model_type, model_num)\n",
-        "model_params = np.load(f\"params/params_model_{model_num}_{model_type}.npz\")\n",
+        "model_params = np.load(f\"params/params_model_{model_num}{model_type}.npz\")\n",
         "w = model_params['alphafold/alphafold_iteration/evoformer/single_activations//weights']\n",
         "b = model_params['alphafold/alphafold_iteration/evoformer/single_activations//bias']\n",
         "single_rep_256 = data[\"representations\"][\"single\"]\n",
@@ -73,7 +77,7 @@
         "id": "KXV8XWuuT0pN",
         "outputId": "02c3003e-e72c-4e70-cb8c-7625def016d6"
       },
-      "execution_count": 22,
+      "execution_count": null,
       "outputs": [
         {
           "output_type": "stream",