Skip to content
Merged

sync #278

Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .pre-commit-config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -31,3 +31,7 @@ repos:
language: system
types: [python]
pass_filenames: true # For speed, we only check the files that are changed
- repo: https://github.com/gitleaks/gitleaks
rev: v8.24.2
hooks:
- id: gitleaks
72 changes: 0 additions & 72 deletions esm/sdk/forge.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,14 +29,6 @@
from esm.utils.constants.api import MIMETYPE_ES_PICKLE
from esm.utils.misc import deserialize_tensors, maybe_list, maybe_tensor
from esm.utils.msa import MSA
from esm.utils.structure.input_builder import (
StructurePredictionInput,
serialize_structure_prediction_input,
)
from esm.utils.structure.molecular_complex import (
MolecularComplex,
MolecularComplexResult,
)
from esm.utils.types import FunctionAnnotation


Expand Down Expand Up @@ -217,70 +209,6 @@ def fold(

return self._process_fold_response(data, sequence)

@retry_decorator
async def async_fold_all_atom(
self, all_atom_input: StructurePredictionInput, model_name: str | None = None
) -> MolecularComplexResult | list[MolecularComplexResult] | ESMProteinError:
"""Fold a molecular complex containing proteins, nucleic acids, and/or ligands.

Args:
all_atom_input: StructurePredictionInput containing sequences for different molecule types
model_name: Override the client level model name if needed
"""
request = self._process_fold_all_atom_request(
all_atom_input, model_name if model_name is not None else self.model
)

try:
data = await self._async_post("fold_all_atom", request)
except ESMProteinError as e:
return e

return self._process_fold_all_atom_response(data)

@retry_decorator
def fold_all_atom(
self, all_atom_input: StructurePredictionInput, model_name: str | None = None
) -> MolecularComplexResult | list[MolecularComplexResult] | ESMProteinError:
"""Predict coordinates for a molecular complex containing proteins, dna, rna, and/or ligands.

Args:
all_atom_input: StructurePredictionInput containing sequences for different molecule types
model_name: Override the client level model name if needed
"""
request = self._process_fold_all_atom_request(
all_atom_input, model_name if model_name is not None else self.model
)

try:
data = self._post("fold_all_atom", request)
except ESMProteinError as e:
return e

return self._process_fold_all_atom_response(data)

@staticmethod
def _process_fold_all_atom_request(
all_atom_input: StructurePredictionInput, model_name: str | None = None
) -> dict[str, Any]:
request: dict[str, Any] = {
"all_atom_input": serialize_structure_prediction_input(all_atom_input),
"model": model_name,
}

return request

@staticmethod
def _process_fold_all_atom_response(data: dict[str, Any]) -> MolecularComplexResult:
complex_data = data.get("complex")
molecular_complex = MolecularComplex.from_state_dict(complex_data)
return MolecularComplexResult(
complex=molecular_complex,
plddt=maybe_tensor(data.get("plddt"), convert_none_to_nan=True),
ptm=data.get("ptm", None),
distogram=maybe_tensor(data.get("distogram"), convert_none_to_nan=True),
)

@retry_decorator
async def async_inverse_fold(
self,
Expand Down
Loading
Loading