Remove needless code

src/transformers/models/esm/modeling_esm.py

@@ -30,11 +30,7 @@
     SequenceClassifierOutput,
     TokenClassifierOutput,
 )
-from ...modeling_utils import (
-    PreTrainedModel,
-    find_pruneable_heads_and_indices,
-    prune_linear_layer,
-)
+from ...modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer
 from ...utils import logging
 from .configuration_esm import EsmConfig
 

src/transformers/models/esm/modeling_esmfold.py

@@ -25,18 +25,20 @@
 )
 from .configuration_esm import EsmConfig
 from .modeling_esm import ESM_START_DOCSTRING, EsmModel, EsmPreTrainedModel
-from .openfold_data_transforms import make_atom14_masks
-from .openfold_np import residue_constants
-from .openfold_np.protein import Protein as OFProtein
-from .openfold_np.protein import to_pdb
-from .openfold_utils.chunk_utils import chunk_layer
-from .openfold_utils.feats import (
+from .openfold_utils import (
+    OFProtein,
+    Rigid,
+    Rotation,
     atom14_to_atom37,
+    chunk_layer,
+    compute_predicted_aligned_error,
+    compute_tm,
     frames_and_literature_positions_to_atom14_pos,
+    make_atom14_masks,
+    residue_constants,
+    to_pdb,
     torsion_angles_to_frames,
 )
-from .openfold_utils.loss import compute_predicted_aligned_error, compute_tm
-from .openfold_utils.rigid_utils import Rigid, Rotation
 
 
 logger = logging.get_logger(__name__)

src/transformers/models/esm/openfold_utils/__init__.py

@@ -0,0 +1,8 @@
+# flake8: noqa
+from .chunk_utils import chunk_layer
+from .data_transforms import make_atom14_masks
+from .feats import atom14_to_atom37, frames_and_literature_positions_to_atom14_pos, torsion_angles_to_frames
+from .loss import compute_predicted_aligned_error, compute_tm
+from .protein import Protein as OFProtein
+from .protein import to_pdb
+from .rigid_utils import Rigid, Rotation

src/transformers/models/esm/openfold_data_transforms.py → src/transformers/models/esm/openfold_utils/data_transforms.py

@@ -14,15 +14,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import itertools
-from functools import reduce, wraps
-from operator import add
-
 import numpy as np
 import torch
 
-from .openfold_np import residue_constants as rc
-from .openfold_utils.tensor_utils import tensor_tree_map, tree_map
+from . import residue_constants as rc
+from .tensor_utils import tensor_tree_map, tree_map
 
 
 def make_atom14_masks(protein):

src/transformers/models/esm/openfold_utils/feats.py

@@ -14,15 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
-from typing import Dict
-
-import numpy as np
 import torch
 import torch.nn as nn
 
-from ..openfold_np import protein
-from ..openfold_np import residue_constants as rc
+from . import residue_constants as rc
 from .rigid_utils import Rigid, Rotation
 from .tensor_utils import batched_gather, one_hot, tensor_tree_map, tree_map