Spaces:
Runtime error
Runtime error
import streamlit as st | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
import torch | |
from transformers import AutoConfig | |
model_name = "farhan2206/dnabert2fourth" | |
# Load the tokenizer and model | |
# Load the configuration associated with the model | |
config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) | |
# Load the model using the correct configuration | |
tokenizer = AutoTokenizer.from_pretrained(model_name, config=config, trust_remote_code=True) | |
model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config, trust_remote_code=True) | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
model = model.to(device) | |
# Streamlit UI | |
def main(): | |
st.title("Epigenetic Marks Prediction") | |
st.write("An application of DNA BERT2") | |
# Sidebar with information | |
st.sidebar.header("About") | |
st.sidebar.write("This app uses DNA BERT2 to predict the presence of epigenetic marks in a given DNA sequence.") | |
# User input | |
user_input = st.text_area("Enter a DNA sequence:", height=150) | |
# Predict when the user provides input | |
if st.button("Classify Sequence"): | |
if user_input: | |
# Call the pred function for prediction | |
# predicted_class, confidence = pred(user_input) | |
predicted_class = pred(user_input) | |
# Display the result | |
st.subheader("Prediction Result") | |
if predicted_class == 1: | |
st.success("Epigenetic Mark detected!") | |
else: | |
st.info("No epigenetic mark found.") | |
# # Display progress bars with percentages | |
# st.subheader("Class Distribution") | |
# st.write("1 - Epigenetic mark found") | |
# st.progress(confidence) | |
# st.text(f"{confidence * 100:.2f}%") | |
# st.write("0 - Epigenetic mark not found") | |
# st.progress(1 - confidence) | |
# st.text(f"{(1 - confidence) * 100:.2f}%") | |
else: | |
st.warning("Please enter a DNA sequence for classification.") | |
# Function for prediction | |
# def pred(sequence): | |
# encoded_input = tokenizer(sequence, return_tensors='pt') | |
# # Pass the encoded input through the model | |
# with torch.no_grad(): | |
# outputs = model(input_ids=encoded_input['input_ids'], attention_mask=encoded_input['attention_mask']) | |
# logits = outputs[0] | |
# predicted_class = logits.argmax(-1).item() | |
# confidence = logits.softmax(dim=-1)[0, 1].item() | |
# return predicted_class, confidence | |
def pred(sequence): | |
# Move the input tensors to the GPU | |
encoded_input = tokenizer(sequence, return_tensors='pt').to(device) | |
# Pass the encoded input through the model | |
with torch.no_grad(): | |
outputs = model(input_ids=encoded_input['input_ids'], attention_mask=encoded_input['attention_mask']).to(device) | |
logits = outputs[0] | |
predicted_class = logits.argmax(-1).item() | |
#confidence = logits.softmax(dim=-1)[0, 1].item() | |
return predicted_class | |
#, confidence | |
if __name__ == "__main__": | |
main() | |
# streamlit run app.py --server.port 9000 | |