kiwi-space / app.py
Peter Moc
Update app.py
8979155 verified
import streamlit as st
from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
# Paths to the saved model and tokenizer
model_path = "path_to_save_model"
tokenizer_path = "path_to_save_tokenizer"
# Load the model and tokenizer from the Hugging Face Hub
model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
# Function to predict the sentiment
def predict_sentiment(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
return torch.argmax(probs, dim=-1).item(), probs
# Streamlit interface
st.title("KIWI Classifier")
st.write("Enter a question or statement to classify:")
user_input = st.text_area("Your input", "")
if st.button("Classify"):
if user_input:
label, probabilities = predict_sentiment(user_input)
st.write(f"Prediction: {label}")
st.write(f"Probabilities: {probabilities.tolist()}")
else:
st.write("Please enter some text to classify.")
# Additional instructions or information
st.write("This application uses a fine-tuned BERT model to classify questions and statements.")