File size: 1,195 Bytes
a681e0a a7d2103 ff694a9 a681e0a 436b48d a681e0a a7d2103 2b716d2 a7d2103 a681e0a ff694a9 a681e0a a7d2103 2b716d2 a7d2103 2b716d2 a7d2103 2b716d2 a7d2103 2b716d2 a7d2103 |
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 |
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import streamlit as st
# Set device (GPU if available, otherwise CPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("mavinsao/mi-roberta-classification")
model = AutoModelForSequenceClassification.from_pretrained("mavinsao/mi-roberta-classification").to(device)
# Streamlit app
st.title('Mental Illness Prediction')
# Input text area for user input
sentence = st.text_area("Enter the long sentence to predict your mental illness state:")
# Prediction button
if st.button('Predict'):
# Tokenize the input sentence
inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True).to(device)
# Forward pass
with torch.no_grad():
outputs = model(**inputs)
# Get predicted probabilities
probabilities = torch.sigmoid(outputs.logits).squeeze(dim=0)
# Get predicted labels with probability greater than 0.5
predicted_labels = [label for i, label in enumerate(tokenizer.labels) if probabilities[i] > 0.5]
st.write("Predicted labels:", predicted_labels)
|