test / app.py
Nikhil0987's picture
Update app.py
3a0dedc verified
raw
history blame
1.13 kB
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
@st.cache_resource # 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)