oss sync

esm/models/vqvae.py

@@ -419,8 +419,8 @@ def decode(
         # This might be broken for chainbreak tokens? We might align to the chainbreak
         ptm = compute_tm(
             pae_logits,  # type: ignore
-            aa_mask=~special_tokens_mask,  # ty:ignore[unknown-argument]
-            max_bin=self.max_pae_bin,  # ty:ignore[unknown-argument]
+            aa_mask=~special_tokens_mask,
+            max_bin=self.max_pae_bin,
         )
 
         plddt_logits = self.plddt_head(x)

esm/sdk/forge.py

@@ -428,7 +428,7 @@ async def async_inverse_fold(
         self,
         coordinates: torch.Tensor,
         config: InverseFoldingConfig,
-        potential_sequence_of_concern: bool,
+        potential_sequence_of_concern: bool = False,
         sequence: str | None = None,
         model_name: str | None = None,
     ) -> ESMProtein | ESMProteinError:
@@ -464,7 +464,7 @@ def inverse_fold(
         self,
         coordinates: torch.Tensor,
         config: InverseFoldingConfig,
-        potential_sequence_of_concern: bool,
+        potential_sequence_of_concern: bool = False,
         sequence: str | None = None,
         model_name: str | None = None,
     ) -> ESMProtein | ESMProteinError:

esm/utils/msa/msa.py

@@ -82,12 +82,14 @@ def from_a3m(
         entries = []
         deletion_rows: list[np.ndarray] = []
         for header, raw in islice(read_sequences(path), max_sequences):
-            deletion_rows.append(a3m_deletion_counts(raw))
+            deletion_row = a3m_deletion_counts(raw)
+            if deletion_rows and len(deletion_row) != len(deletion_rows[0]):
+                raise ValueError(
+                    "A3M match-column count mismatch. "
+                    f"Expected: {len(deletion_rows[0])}, Received: {len(deletion_row)}"
+                )
+            deletion_rows.append(deletion_row)
             seq = remove_insertions_from_sequence(raw) if remove_insertions else raw
-            if entries:
-                assert (
-                    len(seq) == len(entries[0].sequence)
-                ), f"Sequence length mismatch. Expected: {len(entries[0].sequence)}, Received: {len(seq)}"
             entries.append(FastaEntry(header, seq))
         deletions = (
             np.stack(deletion_rows).astype(np.float32) if deletion_rows else None
@@ -159,7 +161,6 @@ def from_bytes(cls, data: bytes) -> MSA:
         ]
         return cls(entries)
 
-    # TODO(jmaccarl): set remove_insertions to True by default here to match other utils
     @classmethod
     def from_sequences(
         cls, sequences: list[str], remove_insertions: bool = False
@@ -177,9 +178,9 @@ def state_dict(self, json_serializable: bool = False) -> dict[str, Any]:
         """Serialize for the Forge wire / storage (mirrors ``ProteinComplex``).
 
         ``deletions`` carries the per-(row, match-column) a3m deletion counts (set by
-        :meth:`from_a3m`) alongside the sequences, so the feature survives even when the
-        default ``remove_insertions`` strips the lowercase insertions out of the
-        sequences. With ``json_serializable=True`` the array is returned as a list.
+        :meth:`from_a3m`) alongside the sequences, so the feature survives when
+        ``remove_insertions=True`` strips lowercase insertions. With
+        ``json_serializable=True`` the array is returned as a list.
         """
         dct: dict[str, Any] = {"sequences": self.sequences}
         if self.deletions is not None:

esm/utils/msa/msa_test.py

@@ -52,8 +52,17 @@ def test_a3m_deletion_counts_vectorized():
         )
 
 
+def test_from_a3m_preserves_insertions_by_default(tmp_path):
+    p = tmp_path / "m.a3m"
+    _write_a3m(p, gz=False)
+    msa = MSA.from_a3m(str(p))
+    assert msa.sequences == ["MKLNT", "MKaaLNT", "M-LNcT"]
+    assert msa.deletions is not None
+    np.testing.assert_array_equal(msa.deletions, _EXPECTED_DELETIONS)
+
+
 def test_from_a3m_records_deletions(tmp_path):
-    msa = _a3m_msa(tmp_path)  # remove_insertions=True (default)
+    msa = _a3m_msa(tmp_path)
     # stored sequences are insertion-stripped (equal length = query length)
     assert msa.sequences == ["MKLNT", "MKLNT", "M-LNT"]
     assert msa.deletions is not None

esm/utils/structure/predicted_aligned_error.py

@@ -48,7 +48,7 @@ def compute_predicted_aligned_error(
     return (probs * bins).sum(dim=-1)
 
 
-@torch.no_grad  # ty:ignore[too-many-positional-arguments]
+@torch.no_grad()
 def compute_tm(logits: torch.Tensor, aa_mask: torch.Tensor, max_bin: float = 31.0):
     square_mask = _compute_pae_masks(aa_mask)
     seqlens = aa_mask.sum(-1, keepdim=True)