Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import TFAutoModelForSequenceClassification | |
from transformers import AutoTokenizer | |
import pandas as pd | |
# title | |
st.title('Ravens 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, 50, 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) | |
# try: | |
if model_type == 'Linear T-Cells (MHC Class I Restriction)': | |
try: | |
model = TFAutoModelForSequenceClassification.from_pretrained('classifier') | |
except tf.errors.InvalidArgumentError: | |
st.warning("We're experiencing server issues. Please try again later!", icon="⚠️") | |
elif model_type == 'Linear T-Cells (MHC Class II Restriction)': | |
try: | |
model = TFAutoModelForSequenceClassification.from_pretrained('classifier2') | |
except tf.errors.InvalidArgumentError: | |
st.warning("We're experiencing server issues. Please try again later!", icon="⚠️") | |
elif model_type == 'Linear B-Cell': | |
try: | |
model = TFAutoModelForSequenceClassification.from_pretrained('bcell') | |
except tf.errors.InvalidArgumentError: | |
st.warning("We're experiencing server issues. Please try again later!", icon="⚠️") | |
# submit button | |
if st.button('Submit'): | |
locations = [] | |
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']) | |
# 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'])) | |
# except InvalidArgumentError: | |
# st.warning("We're experiencing server issues. Please try again later!", icon="⚠️") | |