tianzhechu commited on
Commit
fce31e7
·
1 Parent(s): df2fe71
Files changed (1) hide show
  1. src/streamlit_app.py +9 -2
src/streamlit_app.py CHANGED
@@ -2,10 +2,17 @@ import streamlit as st
2
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
  import os
4
 
 
 
 
 
 
 
 
5
  @st.cache_resource
6
  def load_model():
7
- # Create a local cache directory
8
- cache_dir = "./model_cache"
9
  os.makedirs(cache_dir, exist_ok=True)
10
 
11
  tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small", cache_dir=cache_dir)
 
2
  from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
3
  import os
4
 
5
+ # Disable Streamlit usage stats to avoid permission issues
6
+ os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false"
7
+
8
+ # Use /tmp which is usually writable in containers
9
+ os.environ["HF_HOME"] = "/tmp/huggingface"
10
+ os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
11
+
12
  @st.cache_resource
13
  def load_model():
14
+ # Ensure the cache directory exists
15
+ cache_dir = "/tmp/model_cache"
16
  os.makedirs(cache_dir, exist_ok=True)
17
 
18
  tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small", cache_dir=cache_dir)