Update app.py

app.py

@@ -1,31 +1,76 @@
-from analze import *
-
-app = Flask(__name__)
-
-@app.route('/')
-def home():
-    return render_template('home.html')
-
-@app.route('/upload', methods=['GET', 'POST'])
-def upload_file():
-    if request.method == 'POST':
-        # Check if a file was uploaded
-        if 'file' not in request.files:
-            return render_template('home.html', content='No file uploaded.')
-        file = request.files['file']
-        # Check if the file has a filename
-        if file.filename == '':
-            return render_template('home.html', content='No file selected.')
-        filepath = 'email files/' + file.filename
-        return render_template('home.html',
-                               content=check_file_type(file),
-                               features = get_features(filepath),
-                               pre_content=predict_content(text_feature(filepath)),
-                               pre_tag=predict_html(html_tags_feature(filepath)),
-                               pre_num=predict_num(num_feature(filepath)),
-                               pre_extra=predict_extra(extra_feature(filepath)))
-
-    return render_template('home.html')
-
-if __name__ == '__main__':
-    app.run(host='127.0.0.1',port=5000)
+import torch
+import streamlit as st
+from transformers import AutoTokenizer, OPTForCausalLM
+
+
+@st.cache_resource
+def load_model():
+    tokenizer = AutoTokenizer.from_pretrained("facebook/galactica-30b")
+    model = OPTForCausalLM.from_pretrained("facebook/galactica-30b", device_map='auto', low_cpu_mem_usage=True, torch_dtype=torch.float16)
+    model.gradient_checkpointing_enable()
+    return tokenizer, model
+
+
+st.set_page_config(
+    page_title='BioML-SVM',
+    layout="wide"
+)
+
+with st.spinner("Loading Models and Tokens..."):
+    tokenizer, model = load_model()
+
+with st.form(key='my_form'):
+    col1, col2 = st.columns([10, 1])
+    text_input = col1.text_input(label='Enter the amino sequence')
+    with col2:
+        st.text('')
+        st.text('')
+        submit_button = st.form_submit_button(label='Submit')
+
+    if submit_button:
+        st.session_state['result_done'] = False
+    # input_text = "[START_AMINO]GHMQSITAGQKVISKHKNGRFYQCEVVRLTTETFYEVNFDDGSFSDNLYPEDIVSQDCLQFGPPAEGEVVQVRWTDGQVYGAKFVASHPIQMYQVEFEDGSQLVVKRDDVYTLDEELP[END_AMINO]"
+        with st.spinner('Generating...'):
+            # formatted_text = f"[START_AMINO]{text_input}[END_AMINO]"
+            # formatted_text = f"Here is the sequence: [START_AMINO]{text_input}[END_AMINO]"
+            formatted_text = f"{text_input}"
+            input_ids = tokenizer(formatted_text, return_tensors="pt").input_ids.to("cuda")
+            outputs = model.generate(
+                input_ids=input_ids,
+                max_new_tokens=500
+            )
+            result = tokenizer.decode(outputs[0]).replace(formatted_text, "")
+        st.markdown(result)
+
+        if 'result_done' not in st.session_state or not st.session_state.result_done:
+            st.session_state['result_done'] = True
+            st.session_state['previous_state'] = result
+    else:
+        if 'result_done' in st.session_state and st.session_state.result_done:
+            st.markdown(st.session_state.previous_state)
+
+if 'result_done' in st.session_state and st.session_state.result_done:
+    with st.form(key='ask_more'):
+        col1, col2 = st.columns([10, 1])
+        text_input = col1.text_input(label='Ask more question')
+        with col2:
+            st.text('')
+            st.text('')
+            submit_button = st.form_submit_button(label='Submit')
+
+        if submit_button:
+            with st.spinner('Generating...'):
+                # formatted_text = f"[START_AMINO]{text_input}[END_AMINO]"
+                formatted_text = f"Q:{text_input}\n\nA:\n\n"
+                input_ids = tokenizer(formatted_text, return_tensors="pt").input_ids.to("cuda")
+
+                outputs = model.generate(
+                    input_ids=input_ids,
+                    max_length=len(formatted_text) + 500,
+                    do_sample=True,
+                    top_k=40,
+                    num_beams=1,
+                    num_return_sequences=1
+                )
+                result = tokenizer.decode(outputs[0]).replace(formatted_text, "")
+            st.markdown(result)