Output is constant

[smilenaderi]

Hi I'm trying to run it but I always get the same output:\

[smilenaderi]

import esm

# Load ESM-2 model
model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval()  # disables dropout for deterministic results

AA_LIST = ['A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L', 'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y','X', 'B', 'Z' , ]
# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
main_sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
data = [
    ("main", main_sequence),

]
for i, l in enumerate(main_sequence):
    for aa in AA_LIST:
        data.append((f"{main_sequence[i]}{i}{aa}", main_sequence[:i] + aa + main_sequence[i + 1:]))

data = data[:100]
batch_labels, batch_strs, batch_tokens = batch_converter(data)
batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)

# Extract per-residue representations (on CPU)
with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
token_representations = results["representations"][33]

# Generate per-sequence representations via averaging
# NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.
sequence_representations = []
from tqdm import tqdm
for i, tokens_len in tqdm(enumerate(batch_lens)):
    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))



import torch
from abyssal_pytorch import Abyssal

model = Abyssal()

wt_emb = sequence_representations[0].unsqueeze(0)

# ESM2 embedding for the wildtype amino acid.
for i,rep in enumerate(sequence_representations):
    if i > 0:
        mut_emb = rep.unsqueeze(0)   # ESM2 embedding for the mutated amino acid.

        out = model(wt_emb, mut_emb)
        print(f'''{data[i][0]}\t{out.item()}''')

[smilenaderi]

This is output
100it [00:00, 16960.39it/s]
M0A 0.00906735472381115
M0C 0.009067296981811523
M0D 0.009067350067198277
M0E 0.009067357517778873
M0F 0.009067239239811897
M0G 0.009067406877875328
M0H 0.00906722154468298
M0I 0.009067182429134846
M0K 0.009067310020327568
M0L 0.009067293256521225
M0M 0.009067239239811897
M0N 0.009067325852811337
M0P 0.009067135863006115
M0Q 0.009067275561392307
M0R 0.009067283943295479
M0S 0.009067270904779434
M0T 0.009067206643521786
M0V 0.009067187085747719
M0W 0.009067228063941002
M0Y 0.009067212231457233
M0X 0.009067262522876263
M0B 0.009067262522876263
M0Z 0.009067265316843987
K1A 0.009067150764167309
K1C 0.009067093022167683
K1D 0.009067335166037083
K1E 0.009067394770681858
K1F 0.009067011997103691
K1G 0.009067149832844734
K1H 0.009067146107554436
K1I 0.00906703807413578
K1K 0.009067239239811897
K1L 0.00906679779291153
K1M 0.0090668685734272
K1N 0.009067138656973839
K1P 0.009067010134458542
K1Q 0.00906717125326395
K1R 0.009067307226359844
K1S 0.009067065082490444
K1T 0.009067041799426079
K1V 0.009066859260201454
K1W 0.009067079983651638
K1Y 0.009067129343748093
K1X 0.009067128412425518
K1B 0.009067131206393242
K1Z 0.00906713493168354
T2A 0.009067473001778126
T2C 0.009067410603165627
T2D 0.009067372418940067
T2E 0.009067418053746223
T2F 0.009067387320101261
T2G 0.009067313745617867
T2H 0.009067483246326447
T2I 0.009067334234714508
T2K 0.009067500941455364
T2L 0.009067305363714695
T2M 0.00906718336045742
T2N 0.009067226201295853
T2P 0.009067232720553875
T2Q 0.009067478589713573
T2R 0.00906760897487402
T2S 0.009067120961844921
T2T 0.009067239239811897
T2V 0.009067120030522346
T2W 0.009067459031939507
T2Y 0.00906755868345499
T2X 0.009067408740520477
T2B 0.009067394770681858
T2Z 0.009067396633327007
V3A 0.009067191742360592
V3C 0.009067405946552753

[smilenaderi]

@dohlee could you please check it out?