embeddings

Protein embedding extraction helpers for ESM2 and ProtT5 backbones.

Utils

Example Data

df = pd.DataFrame(
    {
        "sequence": [
            "MKWVTFISLLLLFSSAYSRGVFRRDAHKSEVAHRFKDLGE",
            "GSSHHHHHHSSGLVPRGSHMNNIRRVAILAALVAGQALA",
        ]
    }
)
df.shape

df.head()

run_model_examples = False

ESM2

---

source

get_esm

def get_esm(
    df:DataFrame, # dataframe containing protein sequences
    col:str, # column containing amino acid sequences
    model_name:str='esm2_t33_650M_UR50D', batch_size:int=1, # number of sequences per inference batch
)->DataFrame:

Extract ESM2 embeddings (mean pooled per sequence).

if run_model_examples:
    get_esm(df, "sequence").head()
else:
    print("Skipped get_esm example: set run_model_examples=True to allow model downloads.")

ProtT5 Uniref50

---

source

get_t5

def get_t5(
    df:DataFrame, # dataframe containing protein sequences
    col:str='sequence', # column containing amino acid sequences
)->DataFrame:

Extract ProtT5-XL-uniref50 embeddings from protein sequences in a dataframe.

if run_model_examples:
    get_t5(df, "sequence").head()
else:
    print("Skipped get_t5 example: set run_model_examples=True to allow model downloads.")

ProtT5 BFD

---

source

get_t5_bfd

def get_t5_bfd(
    df:DataFrame, # dataframe containing protein sequences
    col:str='sequence', # column containing amino acid sequences
)->DataFrame:

Extract ProtT5-XL-BFD embeddings from protein sequences in a dataframe.

if run_model_examples:
    get_t5_bfd(df, "sequence").head()
else:
    print("Skipped get_t5_bfd example: set run_model_examples=True to allow model downloads.")