add the func of parsing _ptm model

fastfold/utils/import_weights.py

@@ -127,7 +127,10 @@ def assign(translation_dict, orig_weights):
                 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,14 +540,7 @@ 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",
@@ -582,3 +578,5 @@ def import_jax_weights_(model, npz_path, version="model_1"):
 
     # Set weights
     assign(flat, data)
+
+