ss900371tw commited on
Commit
70e2aaa
·
verified ·
1 Parent(s): 96f2763

Update src/streamlit_app.py

Browse files
Files changed (1) hide show
  1. src/streamlit_app.py +2 -1
src/streamlit_app.py CHANGED
@@ -138,7 +138,8 @@ def load_huggingface_llm(model_id, hf_token):
138
  model_id,
139
  token=hf_token, # 傳遞 Token
140
  torch_dtype=torch.bfloat16, # 建議使用 bfloat16 或 float16 減少記憶體
141
- device_map="auto" # 自動將模型分配到 GPU/CPU
 
142
  )
143
 
144
  # 使用 pipeline 簡化呼叫
 
138
  model_id,
139
  token=hf_token, # 傳遞 Token
140
  torch_dtype=torch.bfloat16, # 建議使用 bfloat16 或 float16 減少記憶體
141
+ device_map="auto" ,# 自動將模型分配到 GPU/CPU
142
+ local_files_only=True
143
  )
144
 
145
  # 使用 pipeline 簡化呼叫