fix off-by-one for molecule ids, and make sure num_molecule_types is passed down to InputFeatureEmbedder

Author: lucidrains

alphafold3_pytorch/alphafold3.py

@@ -2692,7 +2692,7 @@ def __init__(
         dim_single = 384,
         dim_pairwise = 128,
         dim_additional_token_feats = 2,
-        num_molecule_types = 32,
+        num_molecule_types = NUM_MOLECULE_IDS,
         atom_transformer_blocks = 3,
         atom_transformer_heads = 4,
         atom_transformer_kwargs: dict = dict(),
@@ -3141,6 +3141,7 @@ def __init__(
         # input feature embedder
 
         self.input_embedder = InputFeatureEmbedder(
+            num_molecule_types = num_molecule_types,
             dim_atom_inputs = dim_atom_inputs,
             dim_atompair_inputs = dim_atompair_inputs,
             atoms_per_window = atoms_per_window,

alphafold3_pytorch/inputs.py

@@ -56,7 +56,7 @@
 IS_METAL_ION_INDEX = -1
 IS_BIOMOLECULE_INDICES = slice(0, 3)
 
-MOLECULE_GAP_ID = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES) - 1
+MOLECULE_GAP_ID = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES)
 MOLECULE_METAL_ION_ID = MOLECULE_GAP_ID + 1
 NUM_MOLECULE_IDS = len(HUMAN_AMINO_ACIDS) + len(RNA_NUCLEOTIDES) + len(DNA_NUCLEOTIDES) + 2
 

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.2"
+version = "0.2.3"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }