File size: 1,325 Bytes
3a10454 |
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 35 36 37 38 39 40 |
import streamlit as st
import altair as alt
import torch
from transformers import AlbertTokenizer, AlbertForSequenceClassification
# Load pre-trained model and tokenizer
model_name = "albert-base-v2"
tokenizer = AlbertTokenizer.from_pretrained(model_name)
model = AlbertForSequenceClassification.from_pretrained(model_name)
# Define function to classify input text
def classify_text(text):
inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
outputs = model(**inputs)
logits = outputs.logits.detach().numpy()[0]
probabilities = torch.softmax(torch.tensor(logits), dim=0).tolist()
return probabilities
# Set up Streamlit app
st.title("ALBERT Text Classification App")
# Create input box for user to enter text
text_input = st.text_area("Enter text to classify", height=200)
# Classify input text and display results
if st.button("Classify"):
if text_input:
probabilities = classify_text(text_input)
df = pd.DataFrame({
'Label': ['Negative', 'Positive'],
'Probability': probabilities
})
chart = alt.Chart(df).mark_bar().encode(
x='Probability',
y=alt.Y('Label', sort=['Negative', 'Positive'])
)
st.write(chart)
else:
st.write("Please enter some text to classify.")
|