Spaces:
No application file
No application file
import streamlit as st | |
from transformers import AutoTokenizer, AutoModelForSequenceClassification | |
# Load the pre-trained model and tokenizer | |
model_name = "bert-base-uncased" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
# Define the prediction function | |
def classify_text(text): | |
# Tokenize the input text | |
encoded_text = tokenizer(text, truncation=True, padding=True, return_tensors="pt") | |
# Make predictions with the model | |
predictions = model(**encoded_text) | |
pred_labels = predictions.logits.argmax(-1).cpu().numpy() | |
# Get the predicted labels and their corresponding probabilities | |
labels = tokenizer.convert_ids_to_labels(pred_labels) | |
probs = predictions.logits.softmax(-1).cpu().numpy()[:, 1] | |
return labels, probs | |
# Create the Streamlit app | |
st.title("Text Classification App") | |
# Input field for user text | |
user_text = st.text_input("Enter text to classify:") | |
# Predict the classification labels and probabilities | |
if user_text: | |
labels, probs = classify_text(user_text) | |
# Display the classification results | |
st.header("Classification Results:") | |
for label, prob in zip(labels, probs): | |
st.write(f"Label: {label} (Probability: {prob:.3f})") | |