Spaces:
Sleeping
Sleeping
File size: 2,021 Bytes
2190be4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import streamlit as st
from transformers import TFAutoModelForSequenceClassification
from transformers import AutoTokenizer
import pandas as pd
# title
st.title('Raven AI')
# text input with label
sequence = st.text_input('Enter Amino Acid Sequence')
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, 20, 10)
threshold = st.slider('Probability Threshold', 0.0, 1.0, 0.5)
model_checkpoint = "facebook/esm2_t6_8M_UR50D"
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
if model_type == 'Linear T-Cells (MHC Class I Restriction)':
model = TFAutoModelForSequenceClassification.from_pretrained('classifier')
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'):
# run model
locations = []
for i in range(len(sequence) - length):
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'])) |