Bot / src /streamlit_app.py
WillyCodesInit's picture
Update src/streamlit_app.py
2b1f45e verified
raw
history blame contribute delete
931 Bytes
import streamlit as st
from transformers import pipeline
# Set up chatbot pipeline (you can change the model)
chatbot = pipeline("text-generation", model="microsoft/DialoGPT-medium", cache_dir="/app/cache")
st.title("Simple Chatbot")
st.write("Chat with an AI below:")
# Session state to store the chat history
if "history" not in st.session_state:
st.session_state.history = []
# User input
user_input = st.text_input("You:", key="input")
if st.button("Send"):
if user_input:
st.session_state.history.append(f"You: {user_input}")
# Get response from model
response = chatbot(user_input, max_length=100, pad_token_id=50256)[0]["generated_text"]
reply = response.split(user_input)[-1].strip()
st.session_state.history.append(f"Bot: {reply}")
else:
st.warning("Please enter a message.")
# Display conversation
for line in st.session_state.history:
st.write(line)