classification / app.py
Anna-UoC's picture
Create app.py
517802f
raw
history blame
984 Bytes
import streamlit as st
import transformers
from transformers import AutoTokenizer
from transformers import AutoModelForSequenceClassification
from transformers import pipeline
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
st.title("Text Classification with BERT")
text = st.text_input("Enter some text")
if text:
encoded_text = tokenizer.encode_plus(text, max_length=128, padding="max_length", truncation=True, return_tensors="pt")
logits = model(encoded_text["input_ids"], attention_mask=encoded_text["attention_mask"]).logits
pred = logits.argmax().item()
st.write("Prediction:", pred)
#st.text_input('First name')
#classifier = pipeline("text-classification", model=model)
#classifier('st.text_input')
#st.text('Fixed width text')
#x = st.slider('Select a value')
#st.write(x, 'squared is', x * x)