Merge pull request #438 from jnwei/pl_upgrades

Upgrades pl_upgrades to match main branch changes.

openfold/utils/tensor_utils.py

@@ -114,8 +114,7 @@ def tree_map(fn, tree, leaf_type):
     elif isinstance(tree, leaf_type):
         return fn(tree)
     else:
-        print(type(tree))
-        raise ValueError("Not supported")
+        raise ValueError(f"Tree of type {type(tree)} not supported")
 
 
 tensor_tree_map = partial(tree_map, leaf_type=torch.Tensor)

scripts/download_openfold_soloseq_params.sh

 #!/bin/bash
 #
-# Copyright 2021 DeepMind Technologies Limited
+# Copyright 2024 AlQuraishi Laboratory
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,9 +14,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-# Downloads OpenFold parameters.
+# Downloads OpenFold SoloSeq (single sequence model) parameters.
 #
-# Usage: bash download_openfold_params_huggingface.sh /path/to/download/directory
+# Usage: bash download_openfold_soloseq_params.sh /path/to/download/directory
 set -e
 
 if [[ $# -eq 0 ]]; then

tests/test_model.py

@@ -202,4 +202,4 @@ class TestModel(unittest.TestCase):
         out_repro = out_repro["sm"]["positions"][-1]
         out_repro = out_repro.squeeze(0)
 
-        self.assertTrue(torch.max(torch.abs(out_gt - out_repro)) < 1e-3)
+        compare_utils.assert_mean_abs_diff_small(out_gt, out_repro, 1e-3)