|
|
"""Base module for model prediction endpoint query.""" |
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
import logging |
|
|
from abc import ABC, abstractmethod |
|
|
from pathlib import Path |
|
|
from typing import Any |
|
|
|
|
|
from folding_studio_data_models import FoldingModel |
|
|
from pydantic import BaseModel |
|
|
|
|
|
|
|
|
class Query(ABC): |
|
|
"""Interface to define folding job queries.""" |
|
|
|
|
|
MODEL: FoldingModel | None = None |
|
|
|
|
|
@classmethod |
|
|
@abstractmethod |
|
|
def from_file(cls, path: str | Path, **kwargs) -> Query: |
|
|
"""Instantiates a Query object from a file.""" |
|
|
... |
|
|
|
|
|
@classmethod |
|
|
@abstractmethod |
|
|
def from_directory(cls, path: str | Path, **kwargs) -> Query: |
|
|
"""Instantiates a Query object from a directory.""" |
|
|
... |
|
|
|
|
|
@classmethod |
|
|
@abstractmethod |
|
|
def from_protein_sequence(cls, protein: str, **kwargs) -> Query: |
|
|
"""Instantiates a Query object from string representation of a protein.""" |
|
|
... |
|
|
|
|
|
@property |
|
|
@abstractmethod |
|
|
def payload(self) -> dict[str, Any]: |
|
|
"""Returns the payload to be sent in the POST request.""" |
|
|
... |
|
|
|
|
|
@property |
|
|
@abstractmethod |
|
|
def parameters(self) -> BaseModel: |
|
|
"""Parameters of the query.""" |
|
|
... |
|
|
|
|
|
def save_parameters(self, output_dir: Path) -> None: |
|
|
"""Writes the input parameters to a JSON file inside the output directory. |
|
|
|
|
|
Args: |
|
|
output_dir (Path): The directory where the inference parameters JSON file will be saved. |
|
|
""" |
|
|
inference_parameters_path = output_dir / "query_parameters.json" |
|
|
output_dir.mkdir(parents=True, exist_ok=True) |
|
|
with inference_parameters_path.open("w", encoding="utf-8") as f: |
|
|
json.dump(self.parameters.model_dump(mode="json"), f, indent=4) |
|
|
logging.info(f"Input parameters written to {inference_parameters_path}") |
|
|
|
