Update app.py

app.py

@@ -1,6 +1,7 @@
 import streamlit as st
 from transformers import AutoTokenizer, EsmModel
 import torch
+import json
 
 
 model_name = st.selectbox(
@@ -21,6 +22,21 @@ def embed(aa_seq):
     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)