Spaces:
Sleeping
Sleeping
import streamlit as st | |
from streamlit_chat import message | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import random | |
# Title and UI Customization | |
st.set_page_config(page_title="AI Chatbot", page_icon="π€", layout="wide") | |
st.markdown( | |
""" | |
<style> | |
.main {background-color: #f5f7fa;} | |
.stButton button {background-color: #6200EE; color: white; border-radius: 12px;} | |
#chat-container {animation: fadeIn 2s ease-in;} | |
@keyframes fadeIn {from {opacity: 0;} to {opacity: 1;}} | |
</style> | |
""", | |
unsafe_allow_html=True, | |
) | |
# Load Hugging Face Model | |
def load_model(): | |
model_name = "gpt2" # Replace with your preferred Hugging Face model | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForCausalLM.from_pretrained(model_name) | |
return tokenizer, model | |
tokenizer, model = load_model() | |
# Store Chat History | |
if "chat_history" not in st.session_state: | |
st.session_state.chat_history = [] | |
# User Interface | |
def chatbot_ui(): | |
st.markdown( | |
""" | |
<div id="chat-container" style="text-align: center;"> | |
<h1>π AI Chatbot</h1> | |
<p>Ask me anything, and I'll do my best to help you!</p> | |
</div> | |
""", | |
unsafe_allow_html=True, | |
) | |
user_input = st.text_input("Type your question:", "") | |
if st.button("Send") and user_input: | |
generate_response(user_input) | |
# Display Chat History | |
for i, chat in enumerate(st.session_state.chat_history): | |
if chat['role'] == 'user': | |
message(chat['content'], is_user=True, key=f'user_{i}') | |
else: | |
message(chat['content'], key=f'bot_{i}') | |
# Generate Response | |
def generate_response(user_input): | |
inputs = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt") | |
outputs = model.generate(inputs, max_length=200, num_return_sequences=1, do_sample=True, temperature=0.7) | |
bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
# Update Chat History | |
st.session_state.chat_history.append({"role": "user", "content": user_input}) | |
st.session_state.chat_history.append({"role": "bot", "content": bot_reply}) | |
chatbot_ui() | |