openprotein-python 0.4.1__tar.gz → 0.4.2__tar.gz

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: openprotein-python
-Version: 0.4.1
+Version: 0.4.2
 Summary: OpenProtein Python interface.
 Home-page: https://docs.openprotein.ai/
 License: MIT
@@ -16,7 +16,7 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3
 Requires-Dist: pandas (>=1)
-Requires-Dist: pydantic (>=1)
+Requires-Dist: pydantic (>=1.0.0,<2.0.0)
 Requires-Dist: requests (>=2)
 Requires-Dist: tqdm (>=4)
 Description-Content-Type: text/markdown

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/align.py

@@ -1,5 +1,5 @@
 from typing import Iterator, Optional, List, BinaryIO, Literal, Union
-from pydantic import BaseModel, Field, validator, root_validator
+from openprotein.pydantic import BaseModel, Field, validator, root_validator
 from enum import Enum
 from io import BytesIO
 import random

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/data.py

@@ -1,6 +1,6 @@
 import pandas as pd
-import pydantic
-from pydantic import BaseModel
+import openprotein.pydantic as pydantic
+from openprotein.pydantic import BaseModel
 from typing import Optional, List, Union
 from datetime import datetime
 from io import BytesIO
@@ -468,7 +468,7 @@ class DataAPI:
         metadata.sequence_length = len(table["sequence"].values[0])
         return AssayDataset(self.session, metadata)
 
-    def get(self, assay_id: str) -> AssayMetadata:
+    def get(self, assay_id: str, verbose: bool = False) -> AssayMetadata:
         """
         Get an assay dataset by its ID.
 

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/design.py

@@ -9,7 +9,7 @@ from openprotein.jobs import JobType, Job
 
 from openprotein.errors import APIError
 from openprotein.futures import FutureFactory, FutureBase
-from pydantic import BaseModel, Field, validator
+from openprotein.pydantic import BaseModel, Field, validator
 from datetime import datetime
 import re
 

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/embedding.py

@@ -11,7 +11,7 @@ from openprotein.api.poet import (
 )
 from openprotein.futures import FutureBase, FutureFactory
 
-from pydantic import BaseModel, parse_obj_as
+from openprotein.pydantic import BaseModel, parse_obj_as
 import numpy as np
 from typing import Optional, List, Union, Any
 import io
@@ -247,7 +247,7 @@ class EmbeddingResultFuture(MappedAsyncJobFuture, FutureBase):
 
     def get(self, verbose=False) -> List:
         return super().get(verbose=verbose)
-    
+
     @property
     def sequences(self):
         if self._sequences is None:
@@ -305,9 +305,7 @@ def embedding_model_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/embed"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -345,9 +343,7 @@ def embedding_model_logits_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/logits"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -385,9 +381,7 @@ def embedding_model_attn_post(
     """
     endpoint = PATH_PREFIX + f"/models/{model_id}/attn"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -500,9 +494,7 @@ def svd_embed_post(session: APISession, svd_id: str, sequences: List[bytes]) ->
     """
     endpoint = PATH_PREFIX + f"/svd/{svd_id}/embed"
 
-    sequences_unicode = [
-        (s if isinstance(s, str) else s.decode()) for s in sequences 
-    ]
+    sequences_unicode = [(s if isinstance(s, str) else s.decode()) for s in sequences]
     body = {
         "sequences": sequences_unicode,
     }
@@ -715,7 +707,7 @@ class SVDModel(AsyncJobFuture, FutureBase):
         """Get job associated with this SVD model"""
         return job_get(self.session, self.id)
 
-    def get(self):
+    def get(self, verbose: bool = False):
         # overload for AsyncJobFuture
         return self
 
@@ -963,7 +955,7 @@ class PoETModel(OpenProteinModel, EmbBase):
         sequences: List[bytes],
         n_components: int = 1024,
         reduction: Optional[str] = None,
-    ) -> SVDModel:
+    ) -> SVDModel:  # type: ignore
         """
         Fit an SVD on the embedding results of this model. 
 

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/fold.py

@@ -3,7 +3,7 @@ from openprotein.api.jobs import Job, MappedAsyncJobFuture
 import openprotein.config as config
 from openprotein.api.embedding import ModelMetadata
 from openprotein.api.align import validate_msa, MSAFuture
-import pydantic
+import openprotein.pydantic as pydantic
 from typing import Optional, List, Union, Tuple
 from openprotein.futures import FutureBase, FutureFactory
 from abc import ABC, abstractmethod

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/jobs.py

@@ -1,12 +1,12 @@
 # Jobs and job centric flows
 
 
-from typing import List, Union
+from typing import List, Union, Optional
 import concurrent.futures
 import time
 
 import tqdm
-import pydantic
+import openprotein.pydantic as pydantic
 
 from openprotein.base import APISession
 import openprotein.config as config
@@ -105,7 +105,7 @@ class JobsAPI:
             more_recent_than=more_recent_than,
         )
 
-    def get(self, job_id) -> Job:
+    def get(self, job_id: str, verbose: bool = False) -> Job:
         """get Job by ID"""
         return load_job(self.session, job_id)
         # return job_get(self.session, job_id)
@@ -150,7 +150,7 @@ class AsyncJobFuture:
     def cancelled(self):
         return self.job.cancelled()
 
-    def get(self, verbose=False):
+    def get(self, verbose: bool = False):
         raise NotImplementedError()
 
     def wait_until_done(
@@ -176,7 +176,7 @@ class AsyncJobFuture:
     def wait(
         self,
         interval: int = config.POLLING_INTERVAL,
-        timeout: int = None,
+        timeout: Optional[int] = None,
         verbose: bool = False,
     ):
         """
@@ -195,7 +195,7 @@ class AsyncJobFuture:
             self.session, interval=interval, timeout=timeout, verbose=verbose
         )
         self.job = job
-        return self.get(verbose=verbose)
+        return self.get()
 
 
 class StreamingAsyncJobFuture(AsyncJobFuture):

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/poet.py

@@ -1,5 +1,5 @@
 from typing import Iterator, Optional, List, Literal, Dict
-from pydantic import BaseModel, validator
+from openprotein.pydantic import BaseModel, validator
 from io import BytesIO
 import random
 import requests

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/predict.py

@@ -1,5 +1,5 @@
 from typing import Optional, List, Union, Any, Dict, Literal
-from pydantic import BaseModel
+from openprotein.pydantic import BaseModel, root_validator
 
 from openprotein.base import APISession
 from openprotein.api.jobs import AsyncJobFuture
@@ -16,6 +16,29 @@ class SequenceDataset(BaseModel):
     sequences: List[str]
 
 
+class _Prediction(BaseModel):
+    """Prediction details."""
+
+    @root_validator(pre=True)
+    def extract_pred(cls, values):
+        p = values.pop("properties")
+        name = list(p.keys())[0]
+        ymu = p[name]["y_mu"]
+        yvar = p[name]["y_var"]
+        p["name"] = name
+        p["y_mu"] = ymu
+        p["y_var"] = yvar
+
+        values.update(p)
+        return values
+
+    model_id: str
+    model_name: str
+    y_mu: Optional[float] = None
+    y_var: Optional[float] = None
+    name: Optional[str]
+
+
 class Prediction(BaseModel):
     """Prediction details."""
 
@@ -35,6 +58,17 @@ class PredictJobBase(Job):
 class PredictJob(PredictJobBase):
     """Properties about predict job returned via API."""
 
+    @root_validator(pre=True)
+    def extract_pred(cls, values):
+        # Extracting 'predictions' and 'sequences' from the input values
+        v = values.pop("result")
+        preds = [i["predictions"] for i in v]
+        seqs = [i["sequence"] for i in v]
+        values["result"] = [
+            {"sequence": i, "predictions": p} for i, p in zip(seqs, preds)
+        ]
+        return values
+
     class SequencePrediction(BaseModel):
         """Sequence prediction."""
 
@@ -42,7 +76,7 @@ class PredictJob(PredictJobBase):
         predictions: List[Prediction] = []
 
     result: Optional[List[SequencePrediction]] = None
-    job_type: Literal[JobType.workflow_predict] = JobType.workflow_predict
+    job_type: str
 
 
 @register_job_type(JobType.worflow_predict_single_site)
@@ -128,9 +162,9 @@ def _create_predict_job(
 def create_predict_job(
     session: APISession,
     sequences: SequenceDataset,
-    train_job: Any,
+    train_job: Optional[Any] = None,
     model_ids: Optional[List[str]] = None,
-) -> PredictJob:
+) -> FutureBase:
     """
     Creates a predict job with a given set of sequences and a train job.
 
@@ -167,8 +201,9 @@ def create_predict_job(
         model_ids = [model_ids]
     endpoint = "v1/workflow/predict"
     payload = {"sequences": sequences.sequences}
+    train_job_id = train_job.id if train_job is not None else None
     return _create_predict_job(
-        session, endpoint, payload, model_ids=model_ids, train_job_id=train_job.id
+        session, endpoint, payload, model_ids=model_ids, train_job_id=train_job_id
     )
 
 
@@ -177,7 +212,7 @@ def create_predict_single_site(
     sequence: SequenceData,
     train_job: Any,
     model_ids: Optional[List[str]] = None,
-) -> PredictJob:
+) -> FutureBase:
     """
     Creates a predict job for single site mutants with a given sequence and a train job.
 
@@ -318,6 +353,7 @@ class PredictFutureMixin:
 
     session: APISession
     job: PredictJob
+    id: Optional[str] = None
 
     def get_results(
         self, page_size: Optional[int] = None, page_offset: Optional[int] = None
@@ -344,6 +380,7 @@ class PredictFutureMixin:
         HTTPError
             If the GET request does not succeed.
         """
+        assert self.id is not None
         if "single_site" in self.job.job_type:
             return get_single_site_prediction_results(
                 self.session, self.id, page_size, page_offset
@@ -352,7 +389,7 @@ class PredictFutureMixin:
             return get_prediction_results(self.session, self.id, page_size, page_offset)
 
 
-class PredictFuture(PredictFutureMixin, AsyncJobFuture, FutureBase):
+class PredictFuture(PredictFutureMixin, AsyncJobFuture, FutureBase):  # type: ignore
     """Future Job for manipulating results"""
 
     job_type = [JobType.workflow_predict, JobType.worflow_predict_single_site]
@@ -372,24 +409,29 @@ class PredictFuture(PredictFutureMixin, AsyncJobFuture, FutureBase):
         return self.job.job_id
 
     def _fmt_results(self, results):
-        properties = set(results[0].dict()["predictions"][0]["properties"].keys())
+        properties = set(
+            list(i["properties"].keys())[0] for i in results[0].dict()["predictions"]
+        )
         dict_results = {}
         for p in properties:
             dict_results[p] = {}
             for i, r in enumerate(results):
                 s = r.sequence
-                props = r.predictions[0].properties[p]
+                props = [i.properties[p] for i in r.predictions if p in i.properties][0]
                 dict_results[p][s] = {"mean": props["y_mu"], "variance": props["y_var"]}
+        dict_results
         return dict_results
 
     def _fmt_ssp_results(self, results):
-        properties = set(results[0].dict()["predictions"][0]["properties"].keys())
+        properties = set(
+            list(i["properties"].keys())[0] for i in results[0].dict()["predictions"]
+        )
         dict_results = {}
         for p in properties:
             dict_results[p] = {}
             for i, r in enumerate(results):
-                s = f"{r.position+1}{r.amino_acid}"
-                props = r.predictions[0].properties[p]
+                s = s = f"{r.position+1}{r.amino_acid}"
+                props = [i.properties[p] for i in r.predictions if p in i.properties][0]
                 dict_results[p][s] = {"mean": props["y_mu"], "variance": props["y_var"]}
         return dict_results
 
@@ -408,19 +450,21 @@ class PredictFuture(PredictFutureMixin, AsyncJobFuture, FutureBase):
         """
         step = self.page_size
 
-        results = []
+        results: List = []
         num_returned = step
         offset = 0
 
         while num_returned >= step:
             try:
                 response = self.get_results(page_offset=offset, page_size=step)
+                assert isinstance(response.result, list)
                 results += response.result
                 num_returned = len(response.result)
                 offset += num_returned
             except APIError as exc:
                 if verbose:
                     print(f"Failed to get results: {exc}")
+
         if self.job.job_type == JobType.workflow_predict:
             return self._fmt_results(results)
         else:
@@ -444,7 +488,7 @@ class PredictService:
     def create_predict_job(
         self,
         sequences: List,
-        train_job: Any,
+        train_job: Optional[Any] = None,
         model_ids: Optional[List[str]] = None,
     ) -> PredictFuture:
         """
@@ -475,17 +519,18 @@ class PredictService:
         APIError
             If the backend refuses the job (due to sequence length or invalid inputs)
         """
-        if train_job.assaymetadata is not None:
-            if train_job.assaymetadata.sequence_length is not None:
-                if any(
-                    [
-                        train_job.assaymetadata.sequence_length != len(s)
-                        for s in sequences
-                    ]
-                ):
-                    raise InvalidParameterError(
-                        f"Predict sequences length {len(sequences[0])}  != training assaydata ({train_job.assaymetadata.sequence_length})"
-                    )
+        if train_job is not None:
+            if train_job.assaymetadata is not None:
+                if train_job.assaymetadata.sequence_length is not None:
+                    if any(
+                        [
+                            train_job.assaymetadata.sequence_length != len(s)
+                            for s in sequences
+                        ]
+                    ):
+                        raise InvalidParameterError(
+                            f"Predict sequences length {len(sequences[0])}  != training assaydata ({train_job.assaymetadata.sequence_length})"
+                        )
         if not train_job.done():
             print(f"WARNING: training job has status {train_job.status}")
             # raise InvalidParameterError(
@@ -494,7 +539,7 @@ class PredictService:
 
         sequence_dataset = SequenceDataset(sequences=sequences)
         return create_predict_job(
-            self.session, sequence_dataset, train_job, model_ids=model_ids
+            self.session, sequence_dataset, train_job, model_ids=model_ids  # type: ignore
         )
 
     def create_predict_single_site(
@@ -546,5 +591,5 @@ class PredictService:
 
         sequence_dataset = SequenceData(sequence=sequence)
         return create_predict_single_site(
-            self.session, sequence_dataset, train_job, model_ids=model_ids
+            self.session, sequence_dataset, train_job, model_ids=model_ids  # type: ignore
         )

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/api/train.py

@@ -1,7 +1,7 @@
 from typing import Optional, List, Union
-from pydantic import BaseModel
+from openprotein.pydantic import BaseModel
 
-import pydantic
+import openprotein.pydantic as pydantic
 from openprotein.base import APISession
 from openprotein.api.jobs import AsyncJobFuture, Job
 from openprotein.futures import FutureFactory, FutureBase

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/futures.py

@@ -1,14 +1,15 @@
 # Store for Model and Future classes
 from openprotein.jobs import job_get, ResultsParser
-from typing import Optional
+from typing import Optional, Any
 
 
 class FutureBase:
     """Base class for all Future classes.
-    
+
     This class needs to be directly inherited for class discovery."""
+
     # overridden by subclasses
-    job_type = None
+    job_type: Optional[Any] = None
 
     @classmethod
     def get_job_type(cls):
@@ -23,7 +24,9 @@ class FutureFactory:
     """Factory class for creating Future instances based on job_type."""
 
     @staticmethod
-    def create_future(session, job_id:Optional[str] = None, response:dict =None, **kwargs):
+    def create_future(
+        session, job_id: Optional[str] = None, response: Optional[dict] = None, **kwargs
+    ):
         """
         Create and return an instance of the appropriate Future class based on the job type.
 
@@ -36,22 +39,22 @@ class FutureFactory:
         - An instance of the appropriate Future class.
         """
 
-        # parse job 
+        # parse job
         if job_id:
             job = job_get(session, job_id)
         else:
-            if 'job' not in kwargs:
+            if "job" not in kwargs:
                 job = ResultsParser.parse_obj(response)
             else:
                 job = kwargs.pop("job")
-                
+
         # Dynamically discover all subclasses of FutureBase
         future_classes = FutureBase.__subclasses__()
-        kwargs = {k:v for k,v in kwargs.items() if v is not None}
+        kwargs = {k: v for k, v in kwargs.items() if v is not None}
 
         # Find the Future class that matches the job type
         for future_class in future_classes:
             if job.job_type in future_class.get_job_type():
-                return future_class(session=session, job=job, **kwargs)
+                return future_class(session=session, job=job, **kwargs)  # type: ignore
 
         raise ValueError(f"Unsupported job type: {job.job_type}")

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/jobs.py

@@ -1,7 +1,7 @@
 from datetime import datetime
 from typing import Optional, Literal
 import time
-from pydantic import BaseModel, Field
+from openprotein.pydantic import BaseModel, Field
 from openprotein.errors import TimeoutException
 from openprotein.base import APISession
 import openprotein.config as config
@@ -13,7 +13,7 @@ from openprotein.schemas import JobStatus, JobType
 
 class Job(BaseModel):
     status: JobStatus
-    job_id: str
+    job_id: Optional[str]  # must be optional as predict can return None
     # new emb service get doesnt have job_type
     job_type: Optional[Literal[tuple(member.value for member in JobType.__members__.values())]]  # type: ignore
     created_date: Optional[datetime] = None

openprotein_python-0.4.2/openprotein/pydantic.py

@@ -0,0 +1,20 @@
+try:
+    from pydantic.v1 import (
+        BaseModel,
+        Field,
+        ConfigDict,
+        validator,
+        root_validator,
+        parse_obj_as,
+    )
+    import pydantic.v1 as pydantic
+except ImportError:
+    from pydantic import (
+        BaseModel,
+        Field,
+        ConfigDict,
+        validator,
+        root_validator,
+        parse_obj_as,
+    )
+    import pydantic

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/openprotein/schemas.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, ConfigDict
+from openprotein.pydantic import BaseModel, ConfigDict
 from enum import Enum
 
 

{openprotein_python-0.4.1 → openprotein_python-0.4.2}/pyproject.toml

@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "openprotein_python"
 packages = [{include = "openprotein"}]
-version = "0.4.1"
+version = "0.4.2"
 description = "OpenProtein Python interface."
 license = "MIT"
 readme = "README.md"
@@ -15,7 +15,7 @@ classifiers = [
 [tool.poetry.dependencies]
 python = "^3.8"
 requests = ">=2"
-pydantic = ">=1"
+pydantic = "==1.*"
 tqdm = ">=4"
 pandas = ">=1"