File size: 5,936 Bytes
44459bb 01fba1c 44459bb 01fba1c 44459bb 01fba1c 44459bb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 |
"""API simple prediction call wrappers."""
import logging
import warnings
from pathlib import Path
import requests
import typer
from folding_studio_data_models import AF2Parameters, OpenFoldParameters
from folding_studio_data_models.request.folding import FoldingModel
from folding_studio.config import API_URL, REQUEST_TIMEOUT
from folding_studio.utils.data_model import (
PredictRequestCustomFiles,
PredictRequestParams,
)
from folding_studio.utils.file_helpers import partition_template_pdb_from_file
from folding_studio.utils.headers import get_auth_headers
from folding_studio.utils.project_validation import define_project_code_or_raise
def single_job_prediction(
fasta_file: Path,
parameters: AF2Parameters | OpenFoldParameters | None = None,
project_code: str | None = None,
api_key: str | None = None,
*,
ignore_cache: bool = False,
**kwargs,
) -> dict:
"""Make a single job prediction from folding parameters and a FASTA file.
This is a helper function to be called in users scripts.
Args:
fasta_file (Path): Input FASTA file
parameters (AF2Parameters | OpenFoldParameters | None, optional): Job parameters.
For backward compatibility, can be aliased with `af2_parameters`. Defaults to None.
project_code (str | None, optional): Project code under which the jobs are billed.
If None, value is attempted to be read from environment. Defaults to None.
ignore_cache (bool, optional): Force the job submission or not. Defaults to False.
Raises:
ValueError: _description_
typer.Exit: If an error occurs during the API call.
Returns:
dict: API response.
"""
old_parameters = kwargs.get("af2_parameters")
if parameters is None:
if old_parameters is None:
msg = "Argument `parameters` must be specified if deprecated alias `af2_parameters` is not. "
raise ValueError(msg)
else:
warnings.warn(
"Argument 'af2_parameters' is deprecated and will be removed in future release; use 'parameters' instead.",
DeprecationWarning,
stacklevel=2,
)
parameters = old_parameters
elif old_parameters is not None:
raise ValueError("Use either 'parameters' or 'af2_parameters', not both.")
project_code = define_project_code_or_raise(project_code=project_code)
custom_files = PredictRequestCustomFiles(
templates=parameters.custom_templates,
msas=parameters.custom_msas,
initial_guess_files=[parameters.initial_guess_file]
if parameters.initial_guess_file
else None,
templates_masks_files=[parameters.templates_masks_file]
if parameters.templates_masks_file
else None,
)
_ = custom_files.upload(api_key=api_key)
params = parameters.model_dump(mode="json")
pdb_ids, _ = partition_template_pdb_from_file(
custom_templates=parameters.custom_templates
)
folding_model = (
FoldingModel.OPENFOLD
if isinstance(parameters, OpenFoldParameters)
else FoldingModel.AF2
)
params.update(
{
"folding_model": folding_model.value,
"custom_msa_files": custom_files.msas,
"custom_template_ids": list(pdb_ids),
"custom_template_files": custom_files.templates,
"initial_guess_file": custom_files.initial_guess_files[0]
if custom_files.initial_guess_files
else None,
"templates_masks_file": custom_files.templates_masks_files[0]
if custom_files.templates_masks_files
else None,
"ignore_cache": ignore_cache,
}
)
url = API_URL + "predict"
response = requests.post(
url,
data=params,
headers=get_auth_headers(api_key),
files=[("fasta_file", fasta_file.open("rb"))],
params={"project_code": project_code},
timeout=REQUEST_TIMEOUT,
)
response.raise_for_status()
logging.info("Single job successfully submitted.")
response_json = response.json()
return response_json
def simple_prediction(
file: Path,
folding_model: FoldingModel,
params: PredictRequestParams,
custom_files: PredictRequestCustomFiles,
project_code: str | None = None,
) -> dict:
"""Make a simple prediction from a file.
Args:
file (Path): Data source file path.
params (PredictRequestParams): API request parameters.
custom_files (PredictRequestCustomFiles): API request custom files.
project_code (str|None): Project code under which the jobs are billed.
Raises:
typer.Exit: If an error occurs during the API call.
"""
project_code = define_project_code_or_raise(project_code=project_code)
url = API_URL + "predict"
_ = custom_files.upload()
params = params.model_dump(mode="json")
params.update(
{
"folding_model": folding_model.value,
"custom_msa_files": custom_files.msas,
"custom_template_files": custom_files.templates,
"initial_guess_file": custom_files.initial_guess_files[0]
if custom_files.initial_guess_files
else None,
"templates_masks_file": custom_files.templates_masks_files[0]
if custom_files.templates_masks_files
else None,
}
)
response = requests.post(
url,
data=params,
headers=get_auth_headers(),
files=[("fasta_file", file.open("rb"))],
params={"project_code": project_code},
timeout=REQUEST_TIMEOUT,
)
if not response.ok:
print(f"An error occurred: {response.content.decode()}")
raise typer.Exit(code=1)
print("Single job successfully submitted.")
response_json = response.json()
return response_json
|