Spaces:
Sleeping
Sleeping
import streamlit as st | |
import torch | |
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification | |
# Title and Description | |
st.title("Simple DistilBERT Chatbot") | |
st.write("This is a basic chatbot prototype. Ask it something!") | |
# Load Model and Tokenizer | |
# Cache for efficiency | |
def load_model_tokenizer(): | |
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased') | |
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased') | |
return tokenizer, model | |
tokenizer, model = load_model_tokenizer() | |
# User Input | |
user_input = st.text_input("You: ") | |
# Generate Response on Button Click | |
if st.button("Send"): | |
if not user_input: | |
st.warning("Please enter some text.") | |
else: | |
# Preprocess and Generate Response (placeholder) | |
encoded_input = preprocess_input(user_input) | |
outputs = model(**encoded_input) | |
# (TODO) Extract relevant info from outputs | |
bot_response = "I'm still under development, but I understand you said: {}".format(user_input) | |
st.write("Bot: " + bot_response) | |