Touch up OF/AF conversion script

scripts/convert_of_weights_to_jax.py

-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
+# Copyright 2022 AlQuraishi Laboratory
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Converts OpenFold .pt checkpoints into AlphaFold .npz ones, which can then be
+# used to run inference using DeepMind's JAX code.
 import argparse
 
 import numpy as np
@@ -54,7 +68,7 @@ def main(args):
     translation = generate_translation_dict(model, args.config_preset)
     translation = process_translation_dict(translation)
     
-    af_weight_template = data = np.load(args.template_npz_path)
+    af_weight_template = np.load(args.template_npz_path)
     af_weight_template = {k:v for k,v in af_weight_template.items() if k in translation}
     zero = lambda n: n * 0
     af_weight_template = tree_map(zero, af_weight_template, np.ndarray)
@@ -80,7 +94,8 @@ if __name__ == "__main__":
         type=str, 
         default="openfold/resources/params/params_model_1_ptm.npz",
         help="""Path to an AlphaFold checkpoint w/ a superset of the OF
-                checkpoint's parameters"""
+                checkpoint's parameters. params_model_1_ptm.npz always works.
+             """
     )
 
     args = parser.parse_args()