File size: 603 Bytes
21ac434 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
@st.cache(allow_output_mutation=True)
def load_model():
tokenizer = AutoTokenizer.from_pretrained("./my_model")
model = AutoModelForSequenceClassification.from_pretrained("./my_model")
return model, tokenizer
model, tokenizer = load_model()
text_input = st.text_area("Enter text here:")
if st.button("Predict"):
inputs = tokenizer(text_input, return_tensors="pt")
outputs = model(**inputs)
prediction = outputs.logits.argmax(-1).item()
st.write(f"Prediction: {prediction}")
|