app.py

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


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)
    
if st.button('Run Function'):
    embed(aa_seq)

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