Spaces:
Sleeping
Sleeping
File size: 2,001 Bytes
1d614bc 0d568c0 117a3f9 03914f2 1d614bc 117a3f9 03914f2 1d614bc 117a3f9 03914f2 117a3f9 1d614bc 117a3f9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 | import os
import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
# ✅ Use a cache directory that Spaces allows
os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
@st.cache_resource
def load_model():
model_name = "microsoft/DialoGPT-small" # switched from -medium
tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp/hf_cache")
model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="/tmp/hf_cache")
return tokenizer, model
tokenizer, model = load_model()
st.set_page_config(page_title="Chatbot 🤖", page_icon="💬", layout="centered")
st.title("🤖 Hugging Face Chatbot (DialoGPT-small)")
if "chat_history_ids" not in st.session_state:
st.session_state.chat_history_ids = None
if "past_inputs" not in st.session_state:
st.session_state.past_inputs = []
if "generated_responses" not in st.session_state:
st.session_state.generated_responses = []
user_input = st.text_input("You: ", "", key="input")
if st.button("Send") and user_input:
new_input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
if st.session_state.chat_history_ids is not None:
bot_input_ids = torch.cat([st.session_state.chat_history_ids, new_input_ids], dim=-1)
else:
bot_input_ids = new_input_ids
st.session_state.chat_history_ids = model.generate(
bot_input_ids,
max_length=1000,
pad_token_id=tokenizer.eos_token_id
)
bot_output = tokenizer.decode(
st.session_state.chat_history_ids[:, bot_input_ids.shape[-1]:][0],
skip_special_tokens=True
)
st.session_state.past_inputs.append(user_input)
st.session_state.generated_responses.append(bot_output)
if st.session_state.generated_responses:
for i in range(len(st.session_state.generated_responses)):
st.markdown(f"**You:** {st.session_state.past_inputs[i]}")
st.markdown(f"**Bot:** {st.session_state.generated_responses[i]}")
|