add the func of parsing _ptm model

fastfold/utils/import_weights.py

@@ -126,8 +126,11 @@ def assign(translation_dict, orig_weights):
                 print(ref[0].shape)
                 print(weights[0].shape)
                 raise
-
-def get_translation_dict(model, is_multimer: bool = False):
+  
+    
+def import_jax_weights_(model, npz_path, version="model_1", is_multimer: bool = False):
+    data = np.load(npz_path)
+    # translations = get_translation_dict(model, is_multimer=("multimer" in version))
     #######################
     # Some templates
     #######################
@@ -537,23 +540,16 @@ def get_translation_dict(model, is_multimer: bool = False):
             },
         }
 
-    return translations
-
-
-
-def import_jax_weights_(model, npz_path, version="model_1"):
-    data = np.load(npz_path)
-
-    translations = get_translation_dict(model, is_multimer=("multimer" in version))
+    # return translations
 
     no_templ = [
-        "model_3",
-        "model_4",
-        "model_5",
-        "model_3_ptm",
-        "model_4_ptm",
-        "model_5_ptm",
-    ]
+            "model_3",
+            "model_4",
+            "model_5",
+            "model_3_ptm",
+            "model_4_ptm",
+            "model_5_ptm",
+        ]
     if version in no_templ:
         evo_dict = translations["evoformer"]
         keys = list(evo_dict.keys())
@@ -582,3 +578,5 @@ def import_jax_weights_(model, npz_path, version="model_1"):
 
     # Set weights
     assign(flat, data)
+
+