new feature

app.py

@@ -13,12 +13,7 @@ model_type = st.radio(
     "Choose Linear Epitope Classifier",
     ('Linear T-Cells (MHC Class I Restriction)', 'Linear T-Cells (MHC Class II Restriction)', 'Linear B-Cell'))
 
-# windows length slider
-length = st.slider('Window Length', 1, 50, 10)
-threshold = st.slider('Probability Threshold', 0.0, 1.0, 0.75)
-
 model_checkpoint = "facebook/esm2_t6_8M_UR50D"
-
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
 
 if model_type == 'Linear T-Cells (MHC Class I Restriction)':
@@ -27,18 +22,20 @@ elif model_type == 'Linear T-Cells (MHC Class II Restriction)':
     model = TFAutoModelForSequenceClassification.from_pretrained('classifier2')
 elif model_type == 'Linear B-Cell':
     model = TFAutoModelForSequenceClassification.from_pretrained('bcell')
-# submit button
-if st.button('Submit'):
-    if(length > len(sequence)):
-        st.write("Please make sure that your window length is less than the sequence length!")
-    else:
-        # run model
+
+propep = st.radio(
+    "Scan over an entire protein or run a peptide sequence?",
+    ('Protein', 'Peptide'))
+
+if propep == 'Peptide':
+    threshold = st.slider('Probability Threshold', 0.0, 1.0, 0.75)
+    
+    if st.button('Submit'):
         locations = []
-        for i in range(len(sequence) - length + 1):
-            peptide_name = sequence[i:i+length]
-            peptide = tokenizer(peptide_name, return_tensors="tf")
-            output = model(peptide)
-            locations.append([peptide_name, output.logits.numpy()[0][0]])
+        peptide_name = sequence
+        peptide = tokenizer(peptide_name, return_tensors="tf")
+        output = model(peptide)
+        locations.append([peptide_name, output.logits.numpy()[0][0]])
         
         locations = pd.DataFrame(locations, columns = ['Peptide', 'Probability'])
 
@@ -62,4 +59,47 @@ if st.button('Submit'):
 
             return f'background-color: rgb({r}, {g}, {b})'
 
-        st.table(locations.style.applymap(color_survived, subset=['Probability']))
+        st.table(locations.style.applymap(color_survived, subset=['Probability']))
+    
+
+elif propep == 'Protein':
+    # windows length slider
+    length = st.slider('Window Length', 1, 50, 10)
+    threshold = st.slider('Probability Threshold', 0.0, 1.0, 0.75)
+
+    # submit button
+    if st.button('Submit'):
+        if(length > len(sequence)):
+            st.write("Please make sure that your window length is less than the sequence length!")
+        else:
+            # run model
+            locations = []
+            for i in range(len(sequence) - length + 1):
+                peptide_name = sequence[i:i+length]
+                peptide = tokenizer(peptide_name, return_tensors="tf")
+                output = model(peptide)
+                locations.append([peptide_name, output.logits.numpy()[0][0]])
+            
+            locations = pd.DataFrame(locations, columns = ['Peptide', 'Probability'])
+
+            # display table with sequence and probability as the headers
+            def color_survived(x: float): # x between 0 and 1
+                # red to green scale based on x
+                # 0 -> red
+                # 0.5 -> clear
+                # 1 -> green
+
+                # red
+                if x < threshold:
+                    r = 179
+                    g = 40
+                    b = 2
+                # green
+                else:
+                    r = 18
+                    g = 150
+                    b = 6
+
+                return f'background-color: rgb({r}, {g}, {b})'
+
+            st.table(locations.style.applymap(color_survived, subset=['Probability']))