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']))