biolmai 0.1.8__tar.gz → 0.1.10__tar.gz

{biolmai-0.1.8 → biolmai-0.1.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: biolmai
-Version: 0.1.8
+Version: 0.1.10
 Summary: Python client and SDK for https://biolm.ai
 Home-page: https://github.com/BioLM/py-biolm
 Author: Nikhil Haas
@@ -18,6 +18,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Requires-Python: >=3.6
+Provides-Extra: aiodns
 License-File: LICENSE
 License-File: AUTHORS.rst
 

{biolmai-0.1.8 → biolmai-0.1.10}/biolmai/__init__.py

@@ -1,7 +1,8 @@
 """Top-level package for BioLM AI."""
 __author__ = """Nikhil Haas"""
 __email__ = "nikhil@biolm.ai"
-__version__ = '0.1.8'
+__version__ = '0.1.10'
 
+from biolmai.cls import *
 
 __all__ = []

{biolmai-0.1.8 → biolmai-0.1.10}/biolmai/api.py

@@ -15,7 +15,7 @@ import biolmai.auth
 from biolmai.asynch import async_api_call_wrapper
 from biolmai.biolmai import log
 from biolmai.const import MULTIPROCESS_THREADS
-from biolmai.payloads import INST_DAT_TXT, predict_resp_many_in_one_to_many_singles
+from biolmai.payloads import INST_DAT_TXT, PARAMS_ITEMS, predict_resp_many_in_one_to_many_singles
 
 
 @lru_cache(maxsize=64)
@@ -35,65 +35,82 @@ def text_validator(text, c):
     except Exception as e:
         return str(e)
 
+def combine_validation(x, y):
+    if x is None and y is None:
+        return None
+    elif isinstance(x, str) and y is None:
+        return x
+    elif x is None and isinstance(y, str):
+        return y
+    elif isinstance(x, str) and isinstance(y, str):
+        return f"{x}\n{y}"
+
+
+def validate_action(action):
+    def validate(f):
+        def wrapper(*args, **kwargs):
+            # Get class instance at runtime, so you can access not just
+            # APIEndpoints, but any *parent* classes of that,
+            # like ESMFoldSinglechain.
+            class_obj_self = args[0]
+            try:
+                is_method = inspect.getfullargspec(f)[0][0] == "self"
+            except Exception:
+                is_method = False
 
-def validate(f):
-    def wrapper(*args, **kwargs):
-        # Get class instance at runtime, so you can access not just
-        # APIEndpoints, but any *parent* classes of that,
-        # like ESMFoldSinglechain.
-        class_obj_self = args[0]
-        try:
-            is_method = inspect.getfullargspec(f)[0][0] == "self"
-        except Exception:
-            is_method = False
-
-        # Is the function we decorated a class method?
-        if is_method:
-            name = f"{f.__module__}.{class_obj_self.__class__.__name__}.{f.__name__}"
-        else:
-            name = f"{f.__module__}.{f.__name__}"
-
-        if is_method:
-            # Splits name, e.g. 'biolmai.api.ESMFoldSingleChain.predict'
-            action_method_name = name.split(".")[-1]
-            validate_endpoint_action(
-                class_obj_self.action_class_strings,
-                action_method_name,
-                class_obj_self.__class__.__name__,
-            )
-
-        input_data = args[1]
-        # Validate each row's text/input based on class attribute `seq_classes`
-        for c in class_obj_self.seq_classes:
-            # Validate input data against regex
-            if class_obj_self.multiprocess_threads:
-                validation = input_data.text.apply(text_validator, args=(c,))
+            # Is the function we decorated a class method?
+            if is_method:
+                name = f"{f.__module__}.{class_obj_self.__class__.__name__}.{f.__name__}"
             else:
-                validation = input_data.text.apply(text_validator, args=(c,))
-            if "validation" not in input_data.columns:
-                input_data["validation"] = validation
-            else:
-                input_data["validation"] = input_data["validation"].str.cat(
-                    validation, sep="\n", na_rep=""
+                name = f"{f.__module__}.{f.__name__}"
+
+            if is_method:
+                # Splits name, e.g. 'biolmai.api.ESMFoldSingleChain.predict'
+                action_method_name = name.split(".")[-1]
+                validate_endpoint_action(
+                    class_obj_self.action_class_strings,
+                    action_method_name,
+                    class_obj_self.__class__.__name__,
                 )
 
-        # Mark your batches, excluding invalid rows
-        valid_dat = input_data.loc[input_data.validation.isnull(), :].copy()
-        N = class_obj_self.batch_size  # N rows will go per API request
-        # JOIN back, which is by index
-        if valid_dat.shape[0] != input_data.shape[0]:
-            valid_dat["batch"] = np.arange(valid_dat.shape[0]) // N
-            input_data = input_data.merge(
-                valid_dat.batch, left_index=True, right_index=True, how="left"
-            )
-        else:
-            input_data["batch"] = np.arange(input_data.shape[0]) // N
-
-        res = f(class_obj_self, input_data, **kwargs)
-        return res
-
-    return wrapper
+            input_data = args[1]
+            # Validate each row's text/input based on class attribute `seq_classes`
+            if action == "predict":
+                input_classes = class_obj_self.predict_input_classes
+            elif action == "encode":
+                input_classes = class_obj_self.encode_input_classes
+            elif action == "generate":
+                input_classes = class_obj_self.generate_input_classes
+            elif action == "transform":
+                input_classes = class_obj_self.transform_input_classes
+            for c in input_classes:
+                # Validate input data against regex
+                if class_obj_self.multiprocess_threads:
+                    validation = input_data.text.apply(text_validator, args=(c,))
+                else:
+                    validation = input_data.text.apply(text_validator, args=(c,))
+                if "validation" not in input_data.columns:
+                    input_data["validation"] = validation
+                else:
+                    # masking and loc may be more performant option
+                    input_data["validation"] = input_data["validation"].combine(validation, combine_validation)
+
+            # Mark your batches, excluding invalid rows
+            valid_dat = input_data.loc[input_data.validation.isnull(), :].copy()
+            N = class_obj_self.batch_size  # N rows will go per API request
+            # JOIN back, which is by index
+            if valid_dat.shape[0] != input_data.shape[0]:
+                valid_dat["batch"] = np.arange(valid_dat.shape[0]) // N
+                input_data = input_data.merge(
+                    valid_dat.batch, left_index=True, right_index=True, how="left"
+                )
+            else:
+                input_data["batch"] = np.arange(input_data.shape[0]) // N
+            res = f(class_obj_self, input_data, **kwargs)
+            return res
 
+        return wrapper
+    return validate
 
 def convert_input(f):
     def wrapper(*args, **kwargs):
@@ -123,7 +140,20 @@ def convert_input(f):
 
 
 class APIEndpoint:
-    batch_size = 3  # Overwrite in parent classes as needed
+     # Overwrite in parent classes as needed
+    batch_size = 3
+    params = None
+    action_classes = ()
+    api_version = 2
+
+    predict_input_key = "sequence"
+    encode_input_key = "sequence"
+    generate_input_key = "context"
+
+    predict_input_classes = ()
+    encode_input_classes = ()
+    generate_input_classes = ()
+    transform_input_classes = ()
 
     def __init__(self, multiprocess_threads=None):
         # Check for instance-specific threads, otherwise read from env var
@@ -137,7 +167,7 @@ class APIEndpoint:
             [c.__name__.replace("Action", "").lower() for c in self.action_classes]
         )
 
-    def post_batches(self, dat, slug, action, payload_maker, resp_key):
+    def post_batches(self, dat, slug, action, payload_maker, resp_key, key="sequence", params=None):
         keep_batches = dat.loc[~dat.batch.isnull(), ["text", "batch"]]
         if keep_batches.shape[0] == 0:
             pass  # Do nothing - we made nice JSON errors to return in the DF
@@ -145,7 +175,7 @@ class APIEndpoint:
             # raise AssertionError(err)
         if keep_batches.shape[0] > 0:
             api_resps = async_api_call_wrapper(
-                keep_batches, slug, action, payload_maker, resp_key
+                keep_batches, slug, action, payload_maker, resp_key, api_version=self.api_version, key=key,  params=params,
             )
             if isinstance(api_resps, pd.DataFrame):
                 batch_res = api_resps.explode("api_resp")  # Should be lists of results
@@ -170,11 +200,11 @@ class APIEndpoint:
             dat["api_resp"] = None
         return dat
 
-    def unpack_local_validations(self, dat):
+    def unpack_local_validations(self, dat, response_key):
         dat.loc[dat.api_resp.isnull(), "api_resp"] = (
             dat.loc[~dat.validation.isnull(), "validation"]
             .apply(
-                predict_resp_many_in_one_to_many_singles, args=(None, None, True, None)
+                predict_resp_many_in_one_to_many_singles, args=(None, None, True, None), response_key=response_key
             )
             .explode()
         )
@@ -182,39 +212,46 @@ class APIEndpoint:
         return dat
 
     @convert_input
-    @validate
-    def predict(self, dat):
-        dat = self.post_batches(dat, self.slug, "predict", INST_DAT_TXT, "predictions")
-        dat = self.unpack_local_validations(dat)
+    @validate_action("predict")
+    def predict(self, dat, params=None):
+        if self.api_version == 1:
+            dat = self.post_batches(dat, self.slug, "predict", INST_DAT_TXT, "predictions")
+            dat = self.unpack_local_validations(dat, "predictions")
+        else:
+            dat = self.post_batches(dat, self.slug, "predict", PARAMS_ITEMS, "results", key=self.predict_input_key, params=params)
+            dat = self.unpack_local_validations(dat,"results")
         return dat.api_resp.replace(np.nan, None).tolist()
 
-    def infer(self, dat):
-        return self.predict(dat)
+    def infer(self, dat, params=None):
+        return self.predict(dat, params)
 
     @convert_input
-    @validate
+    @validate_action("transform")  # api v1 legacy action
     def transform(self, dat):
         dat = self.post_batches(
             dat, self.slug, "transform", INST_DAT_TXT, "predictions"
         )
-        dat = self.unpack_local_validations(dat)
+        dat = self.unpack_local_validations(dat,"predictions")
         return dat.api_resp.replace(np.nan, None).tolist()
 
-    # @convert_input
-    # @validate
-    # def encode(self, dat):
-    #     # NOTE: we defined this for the specific case of ESM2
-    #     # TODO: this will be need again in v2 of API contract
-    #     dat = self.post_batches(dat, self.slug, "transform",
-    #                             INST_DAT_TXT, "embeddings")
-    #     dat = self.unpack_local_validations(dat)
-    #     return dat.api_resp.replace(np.nan, None).tolist()
+    @convert_input
+    @validate_action("encode")
+    def encode(self, dat, params=None):
+
+        dat = self.post_batches(dat, self.slug, "encode", PARAMS_ITEMS, "results", key=self.encode_input_key, params=params)
+        dat = self.unpack_local_validations(dat, "results")
+        return dat.api_resp.replace(np.nan, None).tolist()
 
     @convert_input
-    @validate
-    def generate(self, dat):
-        dat = self.post_batches(dat, self.slug, "generate", INST_DAT_TXT, "generated")
-        dat = self.unpack_local_validations(dat)
+    @validate_action("generate")
+    def generate(self, dat, params=None):
+        if self.api_version == 1:
+            dat = self.post_batches(dat, self.slug, "generate", INST_DAT_TXT, "generated")
+            dat = self.unpack_local_validations(dat, "predictions")
+        else:
+            dat = self.post_batches(dat, self.slug, "generate", PARAMS_ITEMS, "results", key=self.generate_input_key, params=params)
+            dat = self.unpack_local_validations(dat, "results")
+
         return dat.api_resp.replace(np.nan, None).tolist()
 
 
@@ -290,9 +327,9 @@ class TransformAction:
         return "TransformAction"
 
 
-# class EncodeAction:
-#     def __str__(self):
-#         return "EncodeAction"
+class EncodeAction:
+     def __str__(self):
+         return "EncodeAction"
 
 
 class ExplainAction:

{biolmai-0.1.8 → biolmai-0.1.10}/biolmai/asynch.py

@@ -7,7 +7,7 @@ import aiohttp.resolver
 from aiohttp import ClientSession
 
 from biolmai.auth import get_user_auth_header
-from biolmai.const import BASE_API_URL, MULTIPROCESS_THREADS
+from biolmai.const import BASE_API_URL, BASE_API_URL_V1, MULTIPROCESS_THREADS
 
 aiohttp.resolver.DefaultResolver = aiohttp.resolver.AsyncResolver
 
@@ -146,11 +146,14 @@ async def async_main(urls, concurrency) -> list:
     return await get_all(urls, concurrency)
 
 
-async def async_api_calls(model_name, action, headers, payloads, response_key=None):
+async def async_api_calls(model_name, action, headers, payloads, response_key=None, api_version=2):
     """Hit an arbitrary BioLM model inference API."""
     # Normally would POST multiple sequences at once for greater efficiency,
     # but for simplicity sake will do one at at time right now
-    url = f"{BASE_API_URL}/models/{model_name}/{action}/"
+    if api_version == 1:
+        url = f"{BASE_API_URL_V1}/models/{model_name}/{action}/"
+    else:
+        url = f"{BASE_API_URL}/{model_name}/{action}/"
 
     if not isinstance(payloads, (list, dict)):
         err = "API request payload must be a list or dict, got {}"
@@ -180,15 +183,20 @@ async def async_api_calls(model_name, action, headers, payloads, response_key=No
     # return response
 
 
-def async_api_call_wrapper(grouped_df, slug, action, payload_maker, response_key):
+def async_api_call_wrapper(grouped_df, slug, action, payload_maker, response_key, api_version=2, key="sequence", params=None):
     """Wrap API calls to assist with sequence validation as a pre-cursor to
     each API call.
     """
     model_name = slug
     # payload = payload_maker(grouped_df)
-    init_ploads = grouped_df.groupby("batch").apply(
-        payload_maker, include_batch_size=True
-    )
+    if api_version == 1:
+        init_ploads = grouped_df.groupby("batch").apply(
+            payload_maker, include_batch_size=True
+        )
+    else:
+        init_ploads = grouped_df.groupby("batch").apply(
+            payload_maker, key=key, params=params, include_batch_size=True
+        )
     ploads = init_ploads.to_list()
     init_ploads = init_ploads.to_frame(name="pload")
     init_ploads["batch"] = init_ploads.index
@@ -208,7 +216,7 @@ def async_api_call_wrapper(grouped_df, slug, action, payload_maker, response_key
     #     "https://python.org",
     # ]
     # concurrency = 3
-    api_resp = run(async_api_calls(model_name, action, headers, ploads, response_key))
+    api_resp = run(async_api_calls(model_name, action, headers, ploads, response_key, api_version))
     api_resp = [item for sublist in api_resp for item in sublist]
     api_resp = sorted(api_resp, key=lambda x: x["batch_id"])
     # print(api_resp)

biolmai-0.1.10/biolmai/cls.py

@@ -0,0 +1,176 @@
+"""API inference classes."""
+from biolmai.api import APIEndpoint, GenerateAction, PredictAction, TransformAction, EncodeAction
+from biolmai.validate import (AAExtended,
+                              AAExtendedPlusExtra,
+                              AAUnambiguous,
+                              AAUnambiguousPlusExtra,
+                              DNAUnambiguous,
+                              SingleOrMoreOccurrencesOf,
+                              SingleOccurrenceOf,
+                              PDB,
+                              AAUnambiguousEmpty
+                              )
+
+
+class ESMFoldSingleChain(APIEndpoint):
+    slug = "esmfold-singlechain"
+    action_classes = (PredictAction,)
+    predict_input_classes = (AAUnambiguous(),)
+    batch_size = 2
+
+
+class ESMFoldMultiChain(APIEndpoint):
+    slug = "esmfold-multichain"
+    action_classes = (PredictAction,)
+    predict_input_classes = (AAExtendedPlusExtra(extra=[":"]),)
+    batch_size = 2
+
+
+class ESM2(APIEndpoint):
+    """Example.
+
+    .. highlight:: python
+    .. code-block:: python
+
+       {
+         "items": [{
+           "sequence": "MSILVTRPSPAGEELVSRLRTLGQVAWHFPLIEFSPGQQLPQ"
+         }]
+       }
+    """
+
+    action_classes = (EncodeAction, PredictAction, )
+    encode_input_classes = (AAUnambiguous(),)
+    predict_input_classes = (SingleOrMoreOccurrencesOf(token="<mask>"), AAExtendedPlusExtra(extra=["<mask>"]))
+    batch_size = 1
+
+class ESM2_8M(ESM2):
+    slug = "esm2-8m"
+
+class ESM2_35M(ESM2):
+    slug = "esm2-35m"
+
+class ESM2_150M(ESM2):
+    slug = "esm2-150m"
+
+class ESM2_650M(ESM2):
+    slug = "esm2-650m"
+
+class ESM2_3B(ESM2):
+    slug = "esm2-3b"
+
+class ESM1v(APIEndpoint):
+    """Example.
+
+    .. highlight:: python
+    .. code-block:: python
+
+       {
+          "items": [{
+            "sequence": "QERLEUTGR<mask>SLGYNIVAT"
+          }]
+       }
+    """
+    action_classes = (PredictAction,)
+    predict_input_classes = (SingleOccurrenceOf("<mask>"), AAExtendedPlusExtra(extra=["<mask>"]))
+    batch_size = 5
+
+class ESM1v1(ESM1v):
+    slug = "esm1v-n1"
+
+class ESM1v2(ESM1v):
+    slug = "esm1v-n2"
+
+class ESM1v3(ESM1v):
+    slug = "esm1v-n3"
+
+class ESM1v4(ESM1v):
+    slug = "esm1v-n4"
+
+class ESM1v5(ESM1v):
+    slug = "esm1v-n5"
+
+class ESM1vAll(ESM1v):
+    slug = "esm1v-all"
+
+class ESMIF1(APIEndpoint):
+    slug = "esm-if1"
+    action_classes = (GenerateAction,)
+    generate_input_classes = PDB
+    batch_size = 2
+    generate_input_key = "pdb"
+
+
+class ProGen2(APIEndpoint):
+    action_classes = (GenerateAction,)
+    generate_input_classes = (AAUnambiguousEmpty(),)
+    batch_size = 1
+
+class ProGen2Oas(ProGen2):
+    slug = "progen2-oas"
+
+class ProGen2Medium(ProGen2):
+    slug = "progen2-medium"
+
+class ProGen2Large(ProGen2):
+    slug = "progen2-large"
+
+class ProGen2BFD90(ProGen2):
+    slug = "progen2-bfd90"
+
+class AbLang(APIEndpoint):
+    action_classes = (PredictAction, EncodeAction, GenerateAction,)
+    predict_input_classes = (AAUnambiguous(),)
+    encode_input_classes = (AAUnambiguous(),)
+    generate_input_classes = (SingleOrMoreOccurrencesOf(token="*"), AAUnambiguousPlusExtra(extra=["*"]))
+    batch_size = 32
+    generate_input_key = "sequence"
+
+class AbLangHeavy(AbLang):
+    slug = "ablang-heavy"
+
+class AbLangLight(AbLang):
+    slug = "ablang-light"
+
+class DNABERT(APIEndpoint):
+    slug = "dnabert"
+    action_classes = (EncodeAction,)
+    encode_input_classes = (DNAUnambiguous(),)
+    batch_size = 10
+
+class DNABERT2(APIEndpoint):
+    slug = "dnabert2"
+    action_classes = (EncodeAction,)
+    encode_input_classes = (DNAUnambiguous(),)
+    batch_size = 10
+
+class BioLMToxV1(APIEndpoint):
+    """Example.
+
+    .. highlight:: python
+    .. code-block:: python
+
+       {
+         "instances": [{
+           "data": {"text": "MSILVTRPSPAGEELVSRLRTLGQVAWHFPLIEFSPGQQLPQ"}
+         }]
+       }
+    """
+
+    slug = "biolmtox_v1"
+    action_classes = (TransformAction, PredictAction,)
+    predict_input_classes = (AAUnambiguous(),)
+    transform_input_classes = (AAUnambiguous(),)
+    batch_size = 1
+    api_version = 1
+
+class ProteInfer(APIEndpoint):
+    action_classes = (PredictAction,)
+    predict_input_classes = (AAExtended(),)
+    batch_size = 64
+
+class ProteInferEC(ProteInfer):
+    slug = "proteinfer-ec"
+
+class ProteInferGO(ProteInfer):
+    slug = "proteinfer-go"

{biolmai-0.1.8 → biolmai-0.1.10}/biolmai/const.py

@@ -26,4 +26,5 @@ if int(MULTIPROCESS_THREADS) > max_threads or int(MULTIPROCESS_THREADS) > 128:
 elif int(MULTIPROCESS_THREADS) <= 0:
     err = "Environment variable BIOLMAI_THREADS must be a positive integer."
     raise ValueError(err)
-BASE_API_URL = f"{BASE_DOMAIN}/api/v1"
+BASE_API_URL_V1 = f"{BASE_DOMAIN}/api/v1"
+BASE_API_URL = f"{BASE_DOMAIN}/api/v2"

{biolmai-0.1.8 → biolmai-0.1.10}/biolmai/payloads.py

@@ -7,11 +7,22 @@ def INST_DAT_TXT(batch, include_batch_size=False):
         d["batch_size"] = len(d["instances"])
     return d
 
+def PARAMS_ITEMS(batch, key="sequence", params=None, include_batch_size=False):
+    d = {"items": []}
+    for _, row in batch.iterrows():
+        inst = {key: row.text}
+        d["items"].append(inst)
+    if include_batch_size is True:
+        d["batch_size"] = len(d["items"])
+    if isinstance(params, dict):
+        d["params"] = params
+    return d
+
 
 def predict_resp_many_in_one_to_many_singles(
-    resp_json, status_code, batch_id, local_err, batch_size
+    resp_json, status_code, batch_id, local_err, batch_size, response_key = "results"
 ):
-    expected_root_key = "predictions"
+    expected_root_key = response_key
     to_ret = []
     if not local_err and status_code and status_code == 200:
         list_of_individual_seq_results = resp_json[expected_root_key]