TestApp / app.py
menikev's picture
Update app.py
3f4c8fa verified
raw
history blame
635 Bytes
import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
@st.cache(allow_output_mutation=True)
def load_model():
tokenizer = AutoTokenizer.from_pretrained("./prediction_sinhala.ipynb")
model = AutoModelForSequenceClassification.from_pretrained("./prediction_sinhala.ipynb")
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}")