Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,12 +3,12 @@ import streamlit as st
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# ✅ Force cache dir to /tmp (writeable in Hugging Face Spaces)
+# ✅ Use a cache directory that Spaces allows
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
 
 @st.cache_resource
 def load_model():
-    model_name = "microsoft/DialoGPT-medium"
+    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
@@ -16,9 +16,8 @@ def load_model():
 tokenizer, model = load_model()
 
 st.set_page_config(page_title="Chatbot 🤖", page_icon="💬", layout="centered")
-st.title("🤖 Hugging Face Chatbot with Transformers + Streamlit")
+st.title("🤖 Hugging Face Chatbot (DialoGPT-small)")
 
-# Session state
 if "chat_history_ids" not in st.session_state:
     st.session_state.chat_history_ids = None
 if "past_inputs" not in st.session_state:
@@ -26,7 +25,6 @@ if "past_inputs" not in st.session_state:
 if "generated_responses" not in st.session_state:
     st.session_state.generated_responses = []
 
-# User input
 user_input = st.text_input("You: ", "", key="input")
 
 if st.button("Send") and user_input:
@@ -51,7 +49,6 @@ if st.button("Send") and user_input:
     st.session_state.past_inputs.append(user_input)
     st.session_state.generated_responses.append(bot_output)
 
-# Display conversation
 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]}")