| import os | |
| from src.data.esm.sdk.forge import ESM3ForgeInferenceClient | |
| from src.data.esm.utils.forge_context_manager import ForgeBatchExecutor | |
| # Note: please do not import ESM3SageMakerClient here since that requires AWS SDK. | |
| def client( | |
| model="esm3-sm-open-v1", | |
| url="https://forge.evolutionaryscale.ai", | |
| token=os.environ.get("ESM_API_KEY", ""), | |
| request_timeout=None, | |
| ): | |
| """ | |
| Args: | |
| model: Name of the model to use. | |
| url: URL of a forge server. | |
| token: User's API token. | |
| request_timeout: Amount of time to wait for a request to finish. | |
| Default is wait indefinitely. | |
| """ | |
| return ESM3ForgeInferenceClient(model, url, token, request_timeout) | |
| def batch_executor(max_attempts: int = 10): | |
| """ | |
| Args: | |
| max_attempts: Maximum number of attempts to make before giving up. | |
| Usage: | |
| with batch_executor() as executor: | |
| executor.execute_batch(fn, **kwargs) | |
| """ | |
| return ForgeBatchExecutor(max_attempts) | |