NLPquiz / nlpquiz.py
syedabdullah32's picture
Create nlpquiz.py
65bdceb
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})")