openprotein-python 0.8.2__1-py3-none-any.whl

openprotein/embeddings/embeddings.py

@@ -0,0 +1,151 @@
+"""Embeddings API providing the interface for creating embeddings using protein language models."""
+
+from openprotein.base import APISession
+
+from . import api
+from .esm import ESMModel
+from .future import EmbeddingsResultFuture
+from .models import EmbeddingModel
+from .openprotein import OpenProteinModel
+from .poet import PoETModel
+from .poet2 import PoET2Model
+
+
+class EmbeddingsAPI:
+    """
+    Embeddings API providing the interface for creating embeddings using protein language models.
+
+    You can access all our models either via :meth:`get_model` or directly through the session's embedding attribute using the model's ID and the desired method. For example, to use the attention method on the protein sequence model, you would use ``session.embedding.prot_seq.attn()``.
+
+    Examples
+    --------
+    Accessing a model's method:
+
+    .. code-block:: python
+
+        # To call the attention method on the protein sequence model:
+        import openprotein
+        session = openprotein.connect(username="user", password="password")
+        session.embedding.prot_seq.attn()
+
+    Using the `get_model` method:
+
+    .. code-block:: python
+
+        # Get a model instance by name:
+        import openprotein
+        session = openprotein.connect(username="user", password="password")
+        # list available models:
+        print(session.embedding.list_models() )
+        # init model by name
+        model = session.embedding.get_model('prot-seq')
+    """
+
+    # added for static typing, eg pylance, for autocomplete
+    # at init these are all overwritten.
+
+    #: PoET-2 model
+    poet2: PoET2Model
+    #: PoET model
+    poet: PoETModel
+    #: Prot-seq model
+    prot_seq: OpenProteinModel
+    #: Rotaprot model trained on UniRef50
+    rotaprot_large_uniref50w: OpenProteinModel
+    #: Rotaprot model trained on UniRef90
+    rotaprot_large_uniref90_ft: OpenProteinModel
+    poet_2: PoET2Model
+
+    #: ESM1b model
+    esm1b: ESMModel  # alias
+    esm1b_t33_650M_UR50S: ESMModel
+
+    #: ESM1v model
+    esm1v: ESMModel  # alias
+    esm1v_t33_650M_UR90S_1: ESMModel
+    esm1v_t33_650M_UR90S_2: ESMModel
+    esm1v_t33_650M_UR90S_3: ESMModel
+    esm1v_t33_650M_UR90S_4: ESMModel
+    esm1v_t33_650M_UR90S_5: ESMModel
+
+    #: ESM2 model
+    esm2: ESMModel  # alias
+    esm2_t12_35M_UR50D: ESMModel
+    esm2_t30_150M_UR50D: ESMModel
+    esm2_t33_650M_UR50D: ESMModel
+    esm2_t36_3B_UR50D: ESMModel
+    esm2_t6_8M_UR50D: ESMModel
+
+    def __init__(self, session: APISession):
+        self.session = session
+        # dynamically add models from  api list
+        self._load_models()
+
+    def _load_models(self):
+        # Dynamically add model instances as attributes - precludes any drift
+        models = self.list_models()
+        for model in models:
+            model_name = model.id.replace("-", "_")  # hyphens out
+            setattr(self, model_name, model)
+        # Setup aliases safely
+        if getattr(self, "esm1b_t33_650M_UR50S", None):
+            self.esm1b = self.esm1b_t33_650M_UR50S
+        if getattr(self, "esm1v_t33_650M_UR90S_1", None):
+            self.esm1v = self.esm1v_t33_650M_UR90S_1
+        if getattr(self, "esm2_t33_650M_UR50D", None):
+            self.esm2 = self.esm2_t33_650M_UR50D
+        if getattr(self, "poet_2", None):
+            self.poet2 = self.poet_2
+
+    def list_models(self) -> list[EmbeddingModel]:
+        """list models available for creating embeddings of your sequences"""
+        models = []
+        for model_id in api.list_models(self.session):
+            models.append(
+                EmbeddingModel.create(
+                    session=self.session, model_id=model_id, default=EmbeddingModel
+                )
+            )
+        return models
+
+    def get_model(self, name: str) -> EmbeddingModel:
+        """
+        Get model by model_id. 
+
+        ProtembedModel allows all the usual job manipulation: \
+            e.g. making POST and GET requests for this model specifically. 
+
+
+        Parameters
+        ----------
+        model_id : str
+            the model identifier
+
+        Returns
+        -------
+        ProtembedModel
+            The model
+
+        Raises
+        ------
+        HTTPError
+            If the GET request does not succeed.
+        """
+        model_name = name.replace("-", "_")
+        return getattr(self, model_name)
+
+    def __get_results(self, job) -> EmbeddingsResultFuture:
+        """
+        Retrieves the results of an embedding job.
+
+        Parameters
+        ----------
+        job : Job
+            The embedding job whose results are to be retrieved.
+
+        Returns
+        -------
+        EmbeddingResultFuture
+            An instance of EmbeddingResultFuture
+        """
+        return EmbeddingsResultFuture(job=job, session=self.session)

openprotein/embeddings/esm.py

@@ -0,0 +1,33 @@
+"""Community-based ESM models."""
+
+from .models import AttnModel, EmbeddingModel
+
+
+class ESMModel(AttnModel, EmbeddingModel):
+    """
+    Class providing inference endpoints for Facebook's ESM protein language models.
+
+    Examples
+    --------
+    View specific model details (inc supported tokens) with the `?` operator.
+
+    .. code-block:: python
+
+        >>> import openprotein
+        >>> session = openprotein.connect(username="user", password="password")
+        >>> session.embedding.esm2_t12_35M_UR50D?
+    """
+
+    model_id = [
+        "esm1b_t33_650M_UR50S",
+        "esm1v_t33_650M_UR90S_1",
+        "esm1v_t33_650M_UR90S_2",
+        "esm1v_t33_650M_UR90S_3",
+        "esm1v_t33_650M_UR90S_4",
+        "esm1v_t33_650M_UR90S_5",
+        "esm2_t12_35M_UR50D",
+        "esm2_t30_150M_UR50D",
+        "esm2_t33_650M_UR50D",
+        "esm2_t36_3B_UR50D",
+        "esm2_t6_8M_UR50D",
+    ]

openprotein/embeddings/future.py

@@ -0,0 +1,146 @@
+"""Future for embeddings-related jobs."""
+
+from collections import namedtuple
+from typing import Generator
+
+import numpy as np
+
+from openprotein import config
+
+"""Embeddings results represented as futures."""
+
+from openprotein.base import APISession
+from openprotein.jobs import Future, MappedFuture, StreamingFuture
+
+from . import api
+from .schemas import (
+    AttnJob,
+    EmbeddingsJob,
+    GenerateJob,
+    JobType,
+    LogitsJob,
+    ScoreIndelJob,
+    ScoreJob,
+    ScoreSingleSiteJob,
+)
+
+
+class EmbeddingsResultFuture(MappedFuture, Future):
+    """Future for manipulating results for embeddings-related requests."""
+
+    job: EmbeddingsJob | AttnJob | LogitsJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: EmbeddingsJob | AttnJob | LogitsJob,
+        sequences: list[bytes] | list[str] | None = None,
+        max_workers: int = config.MAX_CONCURRENT_WORKERS,
+    ):
+        super().__init__(session=session, job=job, max_workers=max_workers)
+        # ensure all list[bytes]
+        self._sequences = (
+            [seq.encode() if isinstance(seq, str) else seq for seq in sequences]
+            if sequences is not None
+            else sequences
+        )
+
+    def stream(self):
+        return api.request_get_embeddings_stream(session=self.session, job_id=self.id)
+
+    def get(self, verbose=False) -> list[np.ndarray]:
+        return super().get(verbose=verbose)
+
+    @property
+    def sequences(self) -> list[bytes] | list[str]:
+        if self._sequences is None:
+            self._sequences = api.get_request_sequences(
+                session=self.session, job_id=self.job.job_id, job_type=self.job.job_type
+            )
+        return self._sequences
+
+    @property
+    def id(self):
+        return self.job.job_id
+
+    def __keys__(self):
+        """
+        Get the list of sequences submitted for the embed request.
+
+        Returns
+        -------
+        list of bytes
+            List of sequences.
+        """
+        return self.sequences
+
+    def get_item(self, sequence: bytes) -> np.ndarray:
+        """
+        Get embedding results for specified sequence.
+
+        Args:
+            sequence (bytes): sequence to fetch results for
+
+        Returns:
+            np.ndarray: embeddings
+        """
+        data = api.request_get_sequence_result(
+            session=self.session,
+            job_id=self.job.job_id,
+            sequence=sequence,
+            job_type=self.job.job_type,
+        )
+        return api.result_decode(data)
+
+
+class EmbeddingsScoreFuture(StreamingFuture, Future):
+    """Future for manipulating results for embeddings score-related requests."""
+
+    job: ScoreJob | ScoreIndelJob | ScoreSingleSiteJob
+
+    def __init__(
+        self,
+        session: APISession,
+        job: ScoreJob | ScoreSingleSiteJob | GenerateJob,
+        sequences: list[bytes] | list[str] | None = None,
+    ):
+        super().__init__(session=session, job=job)
+        self._sequences = sequences
+
+    @property
+    def sequences(self) -> list[bytes] | list[str]:
+        if self._sequences is None:
+            self._sequences = api.get_request_sequences(self.session, self.job.job_id)
+        return self._sequences
+
+    def stream(self) -> Generator:
+        if self.job_type == JobType.poet_generate:
+            stream = api.request_get_generate_result(
+                session=self.session, job_id=self.id
+            )
+        else:
+            stream = api.request_get_score_result(session=self.session, job_id=self.id)
+        # mut_code, ... for ssp
+        # name, sequence, ... for score
+        header = next(stream)
+        score_start_index = 0
+        for i, col_name in enumerate(header):
+            if col_name.startswith("score"):
+                score_start_index = i
+                break
+        Score = namedtuple("Score", header[:score_start_index] + ["score"])
+        for line in stream:
+            # combine scores into numpy array
+            scores = np.array([float(s) for s in line[score_start_index:]])
+            output = Score(*line[:score_start_index], scores)
+            yield output
+
+
+class EmbeddingsGenerateFuture(EmbeddingsScoreFuture, StreamingFuture, Future):
+    """Future for manipulating results for embeddings generate-related requests."""
+
+    job: GenerateJob
+
+    @property
+    def sequences(self):
+        raise Exception("generate job does not support listing sequences")