Update app.py

app.py

@@ -3,6 +3,13 @@ from transformers import AutoTokenizer, EsmModel
 import torch
 import json
 
+def embed(aa_seq, tokenizer, model):
+    inputs = tokenizer(aa_seq, return_tensors="pt")
+    outputs = model(**inputs)
+    last_hidden_states = outputs.last_hidden_state.to_list()
+    
+    return last_hidden_states
+
 # selecing and loading a model
 model_name = st.selectbox(
     'Choose a model',
@@ -12,36 +19,38 @@ 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')
+#aa_seq_input = 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)
+#uploading AA sequences file
+uploaded_file = st.file_uploader("Upload JSON with AA sequences", type='json')
+if uploaded_file is not None:
+    data = json.load(uploaded_file)
+    st.write(data)
 
-    data = {
-        'aa_seq':aa_seq, 
-        'last_hidden_states':last_hidden_states
-    }
-    json_data = json.dumps(data)
+def embed_upload_file(upload_dict_dania):
+    # upload_dict_dania = {
+    #                    'uid1': ['aa', 'aan'] 
+    #                     }
+    # output = {
+    #          'uid1': {'aa':[[[0.1298, ....]]], 'aan':[[[0.1298, ....]]]} 
+    #          }
+
+    for uid, seqs in upload_json_dania.items():
+        output[uid] = {}
+        for seq in seqs:
+            output[uid][seq] = embed(seq)
+        
+    json_data = json.dumps(output)
 
     st.download_button(
-        label="Download JSON file",
-        data=json_data,
-        file_name="esm-2 last hidden states.json",
-        mime='application/json'
+        label = "Download JSON file",
+        data = json_data,
+        file_name = "esm-2 last hidden states.json",
+        mime = 'application/json'
     )
 
     
 if st.button('Get embedding'):
-    embed(aa_seq)
+    embed_upload_file(data)
 
 st.write('Also, Dania is not gay')
-#uploading AA sequences file
-uploaded_file = st.file_uploader("Choose a JSON file", type='json')
-if uploaded_file is not None:
-    data = json.load(uploaded_file)
-    st.write(data)