a decoder for ESM-2 token representations

[tomerco4]

Hi,
Is there a simple method available for decoding ESM-2 token representations back into the original amino acid sequence?
I need something like this:

model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
batch_converter = alphabet.get_batch_converter()
model.eval() 

data = [
    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE")]

batch_labels, batch_strs, batch_tokens = batch_converter(data)

with torch.no_grad():
    results = model(batch_tokens, repr_layers=[33])

# amino acids representaion
token_representations = results["representations"][33]

original_sequences = **DECODER**(token_representations)

Where 'original_sequences' holds either the two input sequences in 'data' or two sequences that exhibit a high degree of sequence identity to the original ones.
Thanks for the help!

[naailkhan28]

batch_strs contains the original, un-encoded sequences, there are also several ways of converting tokens back to sequences.

You can use [alphabet.get_tok(x) for x in encoded_sequence] to convert a list or Tensor of indices back into a list of amino acids.

Or if you're using HuggingFace you can use Tokenizer.from_pretrained() instead of a batch_converter and then use tokenizer.decode or tokenizer.decode_batch for one sequence or a batch of sequences.

[tomerco4]

Thanks for the super fast response!
I wasn't clear enough about what I want. To clarify, my goal is to decode ESM-2 token representations directly into amino acid sequences, without any prior knowledge of the original sequences that were used to generate them. Therefore, I require a method that can directly convert the token representations back into corresponding amino acid sequences.

[fcharih]

Hello! Did you ever find a way to do this?