Update app.py

app.py

@@ -2,19 +2,28 @@ import streamlit as st
 from transformers import AutoTokenizer, EsmModel
 import torch
 
-model_name = "facebook/esm2_t6_8M_UR50D"
+
+model_name = st.multiselect(
+    'Choose a model',
+    ["facebook/esm2_t6_8M_UR50D", "facebook/esm2_t48_15B_UR50D"],
+    ["facebook/esm2_t6_8M_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')
 
-inputs = tokenizer(aa_seq, return_tensors="pt")
-outputs = model(**inputs)
-
-last_hidden_states = outputs.last_hidden_state
+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('Model:', model_name)
-st.write('Last hidden state shape:', last_hidden_states.shape)
-st.write('Last hidden states:')
-st.write(last_hidden_states)
 st.write('Also, Dania is not gay')