DexterSptizu's picture
Rename app.py.txt to app.py
80ca522 verified
import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
# Load pre-trained model and tokenizer
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
def classify_text(text):
inputs = tokenizer(text, return_tensors="pt")
with torch.no_grad():
logits = model(**inputs).logits
predicted_class_id = logits.argmax().item()
return model.config.id2label[predicted_class_id]
# Streamlit user interface
st.title('Sentiment Analysis with DistilBert')
user_input = st.text_input("Type a sentence to classify", "Hello, my dog is cute")
prediction = classify_text(user_input)
st.write(f'Sentiment: {prediction}')