Update deprecated code (lucidrains#225)

* Update `pkg_resources`

* Update `pydantic` to match v2 requirements

* Bump numpy version to allow >2.0

* Update `typing` to `beartype.typing`

alphafold3_pytorch/alphafold3.py

@@ -21,7 +21,14 @@
     Sequential,
 )
 
-from typing import Callable, Dict, List, Literal, NamedTuple, Tuple
+from beartype.typing import (
+    Callable,
+    Dict,
+    List,
+    Literal,
+    NamedTuple,
+    Tuple,
+)
 
 from alphafold3_pytorch.tensor_typing import (
     Float,

alphafold3_pytorch/attention.py

@@ -1,5 +1,5 @@
 from __future__ import annotations
-from typing import NamedTuple, Tuple
+from beartype.typing import NamedTuple, Tuple
 
 import torch
 from torch import nn, Tensor

alphafold3_pytorch/common/amino_acid_constants.py

@@ -1,6 +1,6 @@
 """Amino acid constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 

alphafold3_pytorch/common/biomolecule.py

@@ -7,7 +7,7 @@
 import random
 from functools import partial
 from types import ModuleType
-from typing import Any, Dict, List, Optional, Set, Tuple
+from beartype.typing import Any, Dict, List, Optional, Set, Tuple
 
 import numpy as np
 from Bio.PDB.mmcifio import MMCIFIO

alphafold3_pytorch/common/dna_constants.py

@@ -1,6 +1,6 @@
 """Deoxyribonucleic acid (DNA) constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 

alphafold3_pytorch/common/ligand_constants.py

@@ -1,6 +1,6 @@
 """Ligand constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 

alphafold3_pytorch/common/mmcif_metadata.py

@@ -1,6 +1,6 @@
 """mmCIF metadata."""
 
-from typing import Mapping, Sequence
+from beartype.typing import Mapping, Sequence
 
 import numpy as np
 

alphafold3_pytorch/common/rna_constants.py

@@ -1,6 +1,6 @@
 """Ribonucleic acid (RNA) constants used in AlphaFold."""
 
-from typing import Final
+from beartype.typing import Final
 
 import numpy as np
 

alphafold3_pytorch/configs.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from alphafold3_pytorch.tensor_typing import typecheck
-from typing import Callable, List, Dict, Literal
+from beartype.typing import Callable, List, Dict, Literal
 
 from alphafold3_pytorch.alphafold3 import Alphafold3
 
@@ -24,7 +24,11 @@
 import yaml
 from pathlib import Path
 
-from pydantic import BaseModel, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    model_validator,
+)
 
 from pydantic.types import (
     FilePath,
@@ -79,9 +83,10 @@ def yaml_config_path_to_dict(
 # base pydantic classes for constructing alphafold3 and trainer from config files
 
 class BaseModelWithExtra(BaseModel):
-    class Config:
-        extra = 'allow'
-        use_enum_values = True
+    model_config = ConfigDict(
+        extra = 'allow',
+        use_enum_values = True,
+    )
 
 class Alphafold3Config(BaseModelWithExtra):
     dim_atom_inputs: int

alphafold3_pytorch/data/data_pipeline.py

@@ -1,7 +1,7 @@
 """General-purpose data pipeline."""
 
 import os
-from typing import Dict, List, MutableMapping, Optional, Tuple
+from beartype.typing import Dict, List, MutableMapping, Optional, Tuple
 
 import numpy as np
 import torch

alphafold3_pytorch/data/kalign.py

@@ -17,7 +17,7 @@
 import subprocess  # nosec
 import tempfile
 from loguru import logger
-from typing import Mapping, Sequence, Tuple
+from beartype.typing import Mapping, Sequence, Tuple
 
 from alphafold3_pytorch.data import msa_parsing
 from alphafold3_pytorch.data.template_parsing import (

alphafold3_pytorch/data/mmcif_parsing.py

@@ -7,7 +7,7 @@
 import logging
 from collections import defaultdict
 from operator import itemgetter
-from typing import Any, Mapping, Optional, Sequence, Set, Tuple
+from beartype.typing import Any, Mapping, Optional, Sequence, Set, Tuple
 
 import numpy as np
 from Bio import PDB

alphafold3_pytorch/data/msa_parsing.py

@@ -6,7 +6,7 @@
 import random
 import re
 import string
-from typing import Literal, Optional, Sequence, Tuple
+from beartype.typing import Literal, Optional, Sequence, Tuple
 
 from alphafold3_pytorch.tensor_typing import typecheck
 

alphafold3_pytorch/data/template_parsing.py

@@ -1,7 +1,7 @@
 import os
 from datetime import datetime
 from loguru import logger
-from typing import Any, Dict, List, Literal, Mapping, Tuple
+from beartype.typing import Any, Dict, List, Literal, Mapping, Tuple
 
 import numpy as np
 import polars as pl

alphafold3_pytorch/data/weighted_pdb_sampler.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import os
-from typing import Dict, Iterator, List, Literal, Tuple
+from beartype.typing import Dict, Iterator, List, Literal, Tuple
 
 import numpy as np
 import polars as pl

alphafold3_pytorch/inputs.py

@@ -14,7 +14,16 @@
 from itertools import groupby
 from pathlib import Path
 from retrying import retry
-from typing import Any, Callable, Dict, List, Literal, Set, Tuple, Type
+from beartype.typing import (
+    Any,
+    Callable,
+    Dict,
+    List,
+    Literal,
+    Set,
+    Tuple,
+    Type,
+)
 
 import einx
 import numpy as np

alphafold3_pytorch/life.py

@@ -1,5 +1,5 @@
 import os
-from typing import Literal
+from beartype.typing import Literal
 
 import gemmi
 import rdkit.Geometry.rdGeometry as rdGeometry

alphafold3_pytorch/trainer.py

@@ -7,7 +7,7 @@
 from alphafold3_pytorch.alphafold3 import Alphafold3
 from alphafold3_pytorch.attention import pad_at_dim, pad_or_slice_to
 
-from typing import TypedDict, List, Callable
+from beartype.typing import TypedDict, List, Callable
 
 from alphafold3_pytorch.tensor_typing import (
     should_typecheck,

alphafold3_pytorch/utils/data_utils.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, List, Literal, Set, Tuple
+from beartype.typing import Any, Dict, List, Literal, Set, Tuple
 
 import numpy as np
 import torch

alphafold3_pytorch/utils/model_utils.py

@@ -1,8 +1,8 @@
 from functools import wraps
-from typing import Callable, List, Tuple, Union
+from beartype.typing import Callable, List, Tuple, Union
 
 import einx
-import pkg_resources
+import importlib.metadata
 import torch
 import torch.nn.functional as F
 from einops import einsum, pack, rearrange, reduce, repeat, unpack
@@ -644,8 +644,9 @@ def package_available(package_name: str) -> bool:
     :return: `True` if the package is available. `False` otherwise.
     """
     try:
-        return pkg_resources.require(package_name) is not None
-    except pkg_resources.DistributionNotFound:
+        importlib.metadata.version(package_name)
+        return True
+    except importlib.metadata.PackageNotFoundError:
         return False
 
 

alphafold3_pytorch/utils/utils.py

@@ -1,6 +1,6 @@
 import numpy as np
 
-from typing import Any, Iterable, List
+from beartype.typing import Any, Iterable, List
 
 
 def exists(val: Any) -> bool:

pyproject.toml

@@ -43,7 +43,7 @@ dependencies = [
     "huggingface_hub>=0.21.4",
     "jaxtyping>=0.2.28",
     "lightning>=2.2.5",
-    "numpy==1.23.5",
+    "numpy>=1.23.5",
     "polars>=1.1.0",
     "pdbeccdutils>=0.8.5",
     "pydantic>=2.8.2",

scripts/cluster_pdb_test_mmcifs.py

@@ -24,7 +24,7 @@
 import os
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Set, Tuple
+from beartype.typing import Dict, List, Set, Tuple
 
 import numpy as np
 import polars as pl

scripts/cluster_pdb_train_mmcifs.py

@@ -25,7 +25,7 @@
 import os
 import subprocess
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Literal, Optional, Set, Tuple, Union
+from beartype.typing import Dict, List, Literal, Optional, Set, Tuple, Union
 
 import numpy as np
 import polars as pl

scripts/cluster_pdb_val_mmcifs.py

@@ -38,7 +38,7 @@
 import subprocess
 from collections import defaultdict
 from concurrent.futures import ProcessPoolExecutor, as_completed
-from typing import Dict, List, Set, Tuple
+from beartype.typing import Dict, List, Set, Tuple
 
 import numpy as np
 import polars as pl