Spaces:
Sleeping
Sleeping
import streamlit as st | |
import torch | |
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification | |
# Load the model and tokenizer | |
# Cache model for efficiency | |
def load_model(): | |
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased') | |
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased') | |
return tokenizer, model | |
tokenizer, model = load_model() | |
# Input/Output areas | |
st.title("Simple Chatbot") | |
user_input = st.text_input("Enter your message:") | |
# Preprocess and generate response when the user hits Enter | |
if user_input: | |
if user_input.lower() == "quit": | |
st.stop() # End the Streamlit app | |
encoded_input = tokenizer(user_input, return_tensors='pt') | |
outputs = model(**encoded_input) | |
logits = outputs.logits | |
predicted_class_id = logits.argmax(-1).item() | |
sentiment_map = {0: "negative", 1: "neutral", 2: "positive"} # Sentiment mapping | |
sentiment = sentiment_map[predicted_class_id] # Get predicted sentiment | |
st.write(f"Predicted Sentiment: {sentiment}") # Display the sentiment |