app.py

import streamlit as st
from transformers import AutoTokenizer, EsmModel
import torch
import json


model_name = st.selectbox(
    'Choose a model',
    ["facebook/esm2_t6_8M_UR50D", "facebook/esm2_t48_15B_UR50D"])

st.write('You selected model:', model_name)

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = EsmModel.from_pretrained(model_name)

aa_seq = st.text_input('Type AA sequance here')

def embed(aa_seq):
    inputs = tokenizer(aa_seq, return_tensors="pt")
    outputs = model(**inputs)
    last_hidden_states = outputs.last_hidden_state
    st.write('Last hidden state shape:', last_hidden_states.shape)
    st.write('Last hidden states:')
    st.write(last_hidden_states)

    data = {
        'aa_seq':aa_seq, 
        'last_hidden_states':last_hidden_states
    }
    json_data = json.dumps(data)


st.download_button(
    label="Download JSON file",
    data=json_data,
    file_name="esm-2 last hidden states.json",
    mime='application/json'
)

    
if st.button('Run Function'):
    embed(aa_seq)

st.write('Also, Dania is not gay')