Update app.py

app.py

@@ -15,10 +15,6 @@ 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_input = st.text_input('Type AA sequance here')
 
 #uploading AA sequences file
@@ -27,7 +23,7 @@ if uploaded_file is not None:
     data = json.load(uploaded_file)
     #st.write(data)
 
-def embed_upload_file(upload_dict_dania):
+def embed_upload_file(upload_dict_dania, tokenizer, model):
     # upload_dict_dania = {
     #                    'uid1': ['aa', 'aan'] 
     #                     }
@@ -38,7 +34,7 @@ def embed_upload_file(upload_dict_dania):
     for uid, seqs in upload_dict_dania.items():
         output[uid] = {}
         for seq in seqs:
-            output[uid][seq] = embed(seq)
+            output[uid][seq] = embed(seq, tokenizer, model)
         
     json_data = json.dumps(output)
 
@@ -51,6 +47,9 @@ def embed_upload_file(upload_dict_dania):
 
     
 if st.button('Get embedding'):
-    embed_upload_file(data)
+    st.write('You selected model:', model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = EsmModel.from_pretrained(model_name)
+    embed_upload_file(data, tokenizer, model)
 
 st.write('Also, Dania is not gay')