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)