biolmai 0.1.4__py2.py3-none-any.whl → 0.1.7__py2.py3-none-any.whl

biolmai/__init__.py

@@ -1,15 +1,7 @@
 """Top-level package for BioLM AI."""
 __author__ = """Nikhil Haas"""
-__email__ = 'nikhil@biolm.ai'
-__version__ = '0.1.4'
+__email__ = "nikhil@biolm.ai"
+__version__ = '0.1.7'
 
-from biolmai.biolmai import get_api_token, api_call
-from biolmai.api import ESMFoldSingleChain, ESMFoldMultiChain
 
-
-__all__ = [
-    "get_api_token",
-    "api_call",
-    "ESMFoldSingleChain",
-    "ESMFoldMultiChain",
-]
+__all__ = []

biolmai/api.py

@@ -1,111 +1,29 @@
 """References to API endpoints."""
-from biolmai import biolmai
+import datetime
 import inspect
-import pandas as pd
-import numpy as np
-from asyncio import create_task, gather, run, sleep
-from biolmai.asynch import async_main, async_api_calls
-
-from biolmai.biolmai import get_user_auth_header
-from biolmai.const import MULTIPROCESS_THREADS
+import time
 from functools import lru_cache
 
-from biolmai.payloads import INST_DAT_TXT
-from biolmai.validate import ExtendedAAPlusExtra, SingleOccurrenceOf, \
-    UnambiguousAA, \
-    UnambiguousAAPlusExtra
-
-
-def predict_resp_many_in_one_to_many_singles(resp_json, status_code,
-                                             batch_id, local_err, batch_size):
-    expected_root_key = 'predictions'
-    to_ret = []
-    if not local_err and status_code and status_code == 200:
-        list_of_individual_seq_results = resp_json[expected_root_key]
-    elif local_err:
-        list_of_individual_seq_results = [{'error': resp_json}]
-    elif status_code and status_code != 200 and isinstance(resp_json, dict):
-        list_of_individual_seq_results = [resp_json] * batch_size
-    else:
-        raise ValueError("Unexpected response in parser")
-    for idx, item in enumerate(list_of_individual_seq_results):
-        d = {'status_code': status_code,
-             'batch_id': batch_id,
-             'batch_item': idx}
-        if not status_code or status_code != 200:
-            d.update(item)  # Put all resp keys at root there
-        else:
-            # We just append one item, mimicking a single seq in POST req/resp
-            d[expected_root_key] = []
-            d[expected_root_key].append(item)
-        to_ret.append(d)
-    return to_ret
-
-
-def async_api_call_wrapper(grouped_df, slug, action, payload_maker,
-                           response_key):
-    """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)
-    ploads = init_ploads.to_list()
-    init_ploads = init_ploads.to_frame(name='pload')
-    init_ploads['batch'] = init_ploads.index
-    init_ploads = init_ploads.reset_index(drop=True)
-    assert len(ploads) == init_ploads.shape[0]
-    for inst, b in zip(ploads, init_ploads['batch'].to_list()):
-        inst['batch'] = b
-
-    headers = get_user_auth_header()  # Need to pull each time
-    urls = [
-        "https://github.com",
-        "https://stackoverflow.com",
-        "https://python.org",
-    ]
-    # concurrency = 3
-    api_resp = run(async_api_calls(model_name, action, headers,
-                                   ploads, response_key))
-    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)
-    # api_resp = biolmai.api_call(model_name, action, headers, payload,
-    #                             response_key)
-    # resp_json = api_resp.json()
-    # batch_id = int(grouped_df.batch.iloc[0])
-    # batch_size = grouped_df.shape[0]
-    # response = predict_resp_many_in_one_to_many_singles(
-    #     resp_json, api_resp.status_code, batch_id, None, batch_size)
-    return api_resp
-
-
-def api_call_wrapper(df, args):
-    """Wrap API calls to assist with sequence validation as a pre-cursor to
-    each API call.
-    """
-    model_name, action, payload_maker, response_key = args
-    payload = payload_maker(df)
-    headers = get_user_auth_header()  # Need to pull each time
-    api_resp = biolmai.api_call(model_name, action, headers, payload,
-                                response_key)
-    resp_json = api_resp.json()
-    batch_id = int(df.batch.iloc[0])
-    batch_size = df.shape[0]
-    response = predict_resp_many_in_one_to_many_singles(
-        resp_json, api_resp.status_code, batch_id, None, batch_size)
-    return response
+import numpy as np
+import pandas as pd
+import requests
+from requests.adapters import HTTPAdapter
+from requests.packages.urllib3.util.retry import Retry
+
+import biolmai
+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
 
 
 @lru_cache(maxsize=64)
 def validate_endpoint_action(allowed_classes, method_name, api_class_name):
-    action_method_name = method_name.split('.')[-1]
+    action_method_name = method_name.split(".")[-1]
     if action_method_name not in allowed_classes:
-        err = 'Only {} supported on {}'
-        err = err.format(
-            list(allowed_classes),
-            api_class_name
-        )
+        err = "Only {} supported on {}"
+        err = err.format(list(allowed_classes), api_class_name)
         raise AssertionError(err)
 
 
@@ -125,24 +43,23 @@ def validate(f):
         # like ESMFoldSinglechain.
         class_obj_self = args[0]
         try:
-            is_method = inspect.getfullargspec(f)[0][0] == 'self'
-        except:
+            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 = '{}.{}.{}'.format(f.__module__, args[0].__class__.__name__,
-                                     f.__name__)
+            name = f"{f.__module__}.{class_obj_self.__class__.__name__}.{f.__name__}"
         else:
-            name = '{}.{}'.format(f.__module__, f.__name__)
+            name = f"{f.__module__}.{f.__name__}"
 
         if is_method:
             # Splits name, e.g. 'biolmai.api.ESMFoldSingleChain.predict'
-            action_method_name = name.split('.')[-1]
+            action_method_name = name.split(".")[-1]
             validate_endpoint_action(
                 class_obj_self.action_class_strings,
                 action_method_name,
-                class_obj_self.__class__.__name__
+                class_obj_self.__class__.__name__,
             )
 
         input_data = args[1]
@@ -150,35 +67,38 @@ def validate(f):
         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, ))
+                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
+                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='')
+                input_data["validation"] = input_data["validation"].str.cat(
+                    validation, sep="\n", na_rep=""
+                )
 
         # 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
+            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')
+                valid_dat.batch, left_index=True, right_index=True, how="left"
+            )
         else:
-            input_data['batch'] = np.arange(input_data.shape[0])//N
+            input_data["batch"] = np.arange(input_data.shape[0]) // N
 
         res = f(class_obj_self, input_data, **kwargs)
         return res
+
     return wrapper
 
 
 def convert_input(f):
     def wrapper(*args, **kwargs):
         # Get the user-input data argument to the decorated function
-        class_obj_self = args[0]
+        # class_obj_self = args[0]
         input_data = args[1]
         # Make sure we have expected input types
         acceptable_inputs = (str, list, tuple, np.ndarray, pd.DataFrame)
@@ -196,12 +116,13 @@ def convert_input(f):
         if isinstance(input_data, pd.DataFrame) and len(input_data.shape) > 1:
             err = "Detected Pandas DataFrame - input a single vector or Series"
             raise AssertionError(err)
-        input_data = pd.DataFrame(input_data, columns=['text'])
+        input_data = pd.DataFrame(input_data, columns=["text"])
         return f(args[0], input_data, **kwargs)
+
     return wrapper
 
 
-class APIEndpoint(object):
+class APIEndpoint:
     batch_size = 3  # Overwrite in parent classes as needed
 
     def __init__(self, multiprocess_threads=None):
@@ -211,32 +132,26 @@ class APIEndpoint(object):
         else:
             self.multiprocess_threads = MULTIPROCESS_THREADS  # Could be False
         # Get correct auth-like headers
-        self.auth_headers = biolmai.get_user_auth_header()
-        self.action_class_strings = tuple([
-            c.__name__.replace('Action', '').lower() for c in self.action_classes
-        ])
+        self.auth_headers = biolmai.auth.get_user_auth_header()
+        self.action_class_strings = tuple(
+            [c.__name__.replace("Action", "").lower() for c in self.action_classes]
+        )
 
-    @convert_input
-    @validate
-    def predict(self, dat):
-        keep_batches = dat.loc[~dat.batch.isnull(), ['text', 'batch']]
+    def post_batches(self, dat, slug, action, payload_maker, resp_key):
+        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
             # err = "No inputs found following local validation"
             # raise AssertionError(err)
         if keep_batches.shape[0] > 0:
             api_resps = async_api_call_wrapper(
-                keep_batches,
-                self.slug,
-                'predict',
-                INST_DAT_TXT,
-                'predictions'
+                keep_batches, slug, action, payload_maker, resp_key
             )
             if isinstance(api_resps, pd.DataFrame):
-                batch_res = api_resps.explode('api_resp')  # Should be lists of results
+                batch_res = api_resps.explode("api_resp")  # Should be lists of results
                 len_res = batch_res.shape[0]
             else:
-                batch_res = pd.DataFrame({'api_resp': api_resps})
+                batch_res = pd.DataFrame({"api_resp": api_resps})
                 len_res = batch_res.shape[0]
             orig_request_rows = keep_batches.shape[0]
             if len_res != orig_request_rows:
@@ -245,150 +160,151 @@ class APIEndpoint(object):
                 raise AssertionError(err)
 
             # Stack the results horizontally w/ original rows of batches
-            keep_batches['prev_idx'] = keep_batches.index
+            keep_batches["prev_idx"] = keep_batches.index
             keep_batches.reset_index(drop=False, inplace=True)
             batch_res.reset_index(drop=True, inplace=True)
-            keep_batches['api_resp'] = batch_res
-            keep_batches.set_index('prev_idx', inplace=True)
-            dat = dat.join(keep_batches.reindex(['api_resp'], axis=1))
+            keep_batches["api_resp"] = batch_res
+            keep_batches.set_index("prev_idx", inplace=True)
+            dat = dat.join(keep_batches.reindex(["api_resp"], axis=1))
         else:
-            dat['api_resp'] = None
+            dat["api_resp"] = None
+        return dat
+
+    def unpack_local_validations(self, dat):
+        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)
+            )
+            .explode()
+        )
 
-        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)).explode()
+        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)
         return dat.api_resp.replace(np.nan, None).tolist()
 
     def infer(self, dat):
         return self.predict(dat)
 
+    @convert_input
     @validate
-    def tokenize(self, dat):
-        payload = {"instances": [{"data": {"text": dat}}]}
-        resp = biolmai.api_call(
-            model_name=self.slug,
-            headers=self.auth_headers,  # From APIEndpoint base class
-            action='transform',
-            payload=payload
+    def transform(self, dat):
+        dat = self.post_batches(
+            dat, self.slug, "transform", INST_DAT_TXT, "predictions"
         )
-        return resp
+        dat = self.unpack_local_validations(dat)
+        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()
 
-class PredictAction(object):
+    @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)
+        return dat.api_resp.replace(np.nan, None).tolist()
 
-    def __str__(self):
-        return 'PredictAction'
 
+def retry_minutes(sess, URL, HEADERS, dat, timeout, mins):
+    """Retry for N minutes."""
+    HEADERS.update({"Content-Type": "application/json"})
+    attempts, max_attempts = 0, 5
+    try:
+        now = datetime.datetime.now()
+        try_until = now + datetime.timedelta(minutes=mins)
+        while datetime.datetime.now() < try_until and attempts < max_attempts:
+            response = None
+            try:
+                log.info(f"Trying {datetime.datetime.now()}")
+                response = sess.post(URL, headers=HEADERS, data=dat, timeout=timeout)
+                if response.status_code not in (400, 404):
+                    response.raise_for_status()
+                if "error" in response.json():
+                    raise ValueError(response.json().dumps())
+                else:
+                    break
+            except Exception as e:
+                log.warning(e)
+                if response:
+                    log.warning(response.text)
+                time.sleep(5)  # Wait 5 seconds between tries
+            attempts += 1
+        if response is None:
+            err = "Got Nonetype response"
+            raise ValueError(err)
+        elif "Server Error" in response.text:
+            err = "Got Server Error"
+            raise ValueError(err)
+    except Exception:
+        return response
+    return response
 
-class GenerateAction(object):
 
+def requests_retry_session(
+    retries=3,
+    backoff_factor=0.3,
+    status_forcelist=None,
+    session=None,
+):
+    if status_forcelist is None:
+        status_forcelist = list(range(400, 599))
+    session = session or requests.Session()
+    retry = Retry(
+        total=retries,
+        read=retries,
+        connect=retries,
+        backoff_factor=backoff_factor,
+        status_forcelist=status_forcelist,
+    )
+    adapter = HTTPAdapter(max_retries=retry)
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
+
+
+class PredictAction:
     def __str__(self):
-        return 'GenerateAction'
-
+        return "PredictAction"
 
-class TransformAction(object):
 
+class GenerateAction:
     def __str__(self):
-        return 'TransformAction'
+        return "GenerateAction"
 
 
-class ExplainAction(object):
-
+class TransformAction:
     def __str__(self):
-        return 'ExplainAction'
+        return "TransformAction"
+
 
+# class EncodeAction:
+#     def __str__(self):
+#         return "EncodeAction"
 
-class SimilarityAction(object):
 
+class ExplainAction:
     def __str__(self):
-        return 'SimilarityAction'
+        return "ExplainAction"
 
 
-class FinetuneAction(object):
+class SimilarityAction:
+    def __str__(self):
+        return "SimilarityAction"
+
 
+class FinetuneAction:
     def __str__(self):
-        return 'FinetuneAction'
-
-
-class ESMFoldSingleChain(APIEndpoint):
-    slug = 'esmfold-singlechain'
-    action_classes = (PredictAction, )
-    seq_classes = (UnambiguousAA(), )
-    batch_size = 2
-
-
-class ESMFoldMultiChain(APIEndpoint):
-    slug = 'esmfold-multichain'
-    action_classes = (PredictAction, )
-    seq_classes = (ExtendedAAPlusExtra(extra=[':']), )
-    batch_size = 2
-
-
-class ESM2Embeddings(APIEndpoint):
-    """Example.
-
-    ```python
-    {
-      "instances": [{
-        "data": {"text": "MSILVTRPSPAGEELVSRLRTLGQVAWHFPLIEFSPGQQLPQ"}
-      }]
-    }
-    ```
-    """
-    slug = 'esm2_t33_650M_UR50D'
-    action_classes = (TransformAction,)
-    seq_classes = (UnambiguousAA(), )
-    batch_size = 3
-
-
-class ESM1v1(APIEndpoint):
-    """Example.
-
-    ```python
-    {
-      "instances": [{
-        "data": {"text": "QERLEUTGR<mask>SLGYNIVAT"}
-      }]
-    }
-    ```
-    """
-    slug = 'esm1v_t33_650M_UR90S_1'
-    action_classes = (PredictAction, )
-    seq_classes = (SingleOccurrenceOf('<mask>'),
-                   ExtendedAAPlusExtra(extra=['<mask>']))
-    batch_size = 5
-
-
-class ESM1v2(APIEndpoint):
-    slug = 'esm1v_t33_650M_UR90S_2'
-    action_classes = (PredictAction, )
-    seq_classes = (SingleOccurrenceOf('<mask>'),
-                   ExtendedAAPlusExtra(extra=['<mask>']))
-    batch_size = 5
-
-
-class ESM1v3(APIEndpoint):
-    slug = 'esm1v_t33_650M_UR90S_3'
-    action_classes = (PredictAction, )
-    seq_classes = (SingleOccurrenceOf('<mask>'),
-                   ExtendedAAPlusExtra(extra=['<mask>']))
-    batch_size = 5
-
-
-class ESM1v4(APIEndpoint):
-    slug = 'esm1v_t33_650M_UR90S_4'
-    action_classes = (PredictAction, )
-    seq_classes = (SingleOccurrenceOf('<mask>'),
-                   ExtendedAAPlusExtra(extra=['<mask>']))
-    batch_size = 5
-
-
-class ESM1v5(APIEndpoint):
-    slug = 'esm1v_t33_650M_UR90S_5'
-    action_classes = (PredictAction, )
-    seq_classes = (SingleOccurrenceOf('<mask>'),
-                   ExtendedAAPlusExtra(extra=['<mask>']))
-    batch_size = 5
+        return "FinetuneAction"