Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,5 +1,7 @@
 import os
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"
+os.environ["HF_HOME"] = "/tmp/hf_cache"
+os.environ["HF_DATASETS_CACHE"] = "/tmp/hf_cache"
 import json
 import requests
 import streamlit as st
@@ -22,8 +24,20 @@ st.write("🔐 토큰 있음:", os.environ.get("HUGGINGFACE_TOKEN") is not None)
 def load_model():
     token = os.environ.get("HUGGINGFACE_TOKEN")
     model_id = "google/gemma-2-2b-it"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
-    model = AutoModelForCausalLM.from_pretrained(model_id, use_auth_token=token)
+    cache_dir = "/tmp/hf_cache"  # ✅ Hugging Face Spaces에서 허용된 디렉토리
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_id,
+        use_auth_token=token,
+        cache_dir=cache_dir
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        use_auth_token=token,
+        cache_dir=cache_dir
+    )
+
     return pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 llm = load_model()