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})") | |