Skip to content

Commit

Permalink
ruff fix
Browse files Browse the repository at this point in the history
  • Loading branch information
shariqm-modal committed Dec 5, 2024
1 parent df7499a commit e92f191
Showing 1 changed file with 15 additions and 8 deletions.
23 changes: 15 additions & 8 deletions 06_gpu_and_ml/protein-folding/boltz1.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,9 +54,10 @@ def main(

output_path = here / "data" / "boltz1_result.tar.gz"
print(f"locally writing output tar to {output_path}")
with open(output_path, 'wb') as f:
with open(output_path, "wb") as f:
f.write(output)


# ## Installing Boltz-1 Python dependencies on Modal

# Code running on Modal runs inside containers built from [container images](https://modal.com/docs/guide/images)
Expand All @@ -69,7 +70,9 @@ def main(

image = (
modal.Image.debian_slim(python_version="3.12")
.run_commands("uv pip install --system --compile-bytecode boltz==0.3.2 hf_transfer==0.1.8")
.run_commands(
"uv pip install --system --compile-bytecode boltz==0.3.2 hf_transfer==0.1.8"
)
.env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

Expand All @@ -89,7 +92,9 @@ def main(
# A Modal Volume is a file system that all of your code running on Modal (or elsewhere!) can access.
# For more on storing model weights on Modal, see [this guide](https://modal.com/docs/guide/model-weights).

boltz_model_volume = modal.Volume.from_name("boltz1-models", create_if_missing=True)
boltz_model_volume = modal.Volume.from_name(
"boltz1-models", create_if_missing=True
)
models_dir = Path("/models/boltz1")

# ## Running Boltz-1 on Modal
Expand All @@ -102,6 +107,7 @@ def main(
# an amino acid sequence as well as a path to an MSA file. We create these
# files paths in the function below so they can be picked up by the inference code.


@app.function(
image=image,
volumes={models_dir: boltz_model_volume},
Expand All @@ -112,17 +118,17 @@ def boltz1_inference(msadata, ligandyaml):
ligand_filename = "ligand.yaml"
temp_filename = "temp.tar.gz"

Path('./seq1.a3m').write_text(msadata)
Path(f'./{ligand_filename}').write_text(ligandyaml)
Path("./seq1.a3m").write_text(msadata)
Path(f"./{ligand_filename}").write_text(ligandyaml)

print(f"predicting using boltz model {models_dir}")
os.system(f'boltz predict {ligand_filename} --cache={models_dir}')
os.system(f"boltz predict {ligand_filename} --cache={models_dir}")

print(f"creating tar file {temp_filename} of outputs")
os.system(f"tar czvf {temp_filename} boltz_results_ligand")

print(f"converting {temp_filename} to bytes for returning")
output = open(f'./{temp_filename}', 'rb').read()
output = open(f"./{temp_filename}", "rb").read()
return output


Expand All @@ -142,6 +148,7 @@ def boltz1_inference(msadata, ligandyaml):
# as a separate task, concurrently. We don't need to worry about this use of `async`
# spreading to the rest of our code -- Modal launches just this Function in an async runtime.


@app.function(
volumes={models_dir: boltz_model_volume},
timeout=20 * MINUTES,
Expand All @@ -152,7 +159,7 @@ async def download_inference_dependencies(force=False):

import aiohttp

base_url = "https://huggingface.co/boltz-community/boltz-1/resolve/e01950840c2a2ec881695f26e994a73b417af0b2/" # sic
base_url = "https://huggingface.co/boltz-community/boltz-1/resolve/e01950840c2a2ec881695f26e994a73b417af0b2/" # sic
inference_dependencies = [
"boltz1.ckpt",
"ccd.pkl",
Expand Down

0 comments on commit e92f191

Please sign in to comment.