Source code for openprotein.embeddings.esm

"""Community-based ESM models."""

from .models import EmbeddingModel




class ESMModel(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",
    ]