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