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 Peptide') | |
| 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')) | |
| 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') | |
| # 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.5) | |
| 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'])) | |
| # 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'])) |