Jay-10020 commited on
Commit
33dfbb4
·
1 Parent(s): 18de33e

fix: use snapshot_download to avoid OOM during docker build

Browse files
Files changed (1) hide show
  1. Dockerfile +5 -4
Dockerfile CHANGED
@@ -46,14 +46,15 @@ RUN pip install --no-cache-dir --user -r requirements.txt
46
  # - Whisper base (~140 MB)
47
  RUN python -c "\
48
  from sentence_transformers import SentenceTransformer; \
49
- from transformers import AutoModelForCausalLM, AutoTokenizer; \
50
- import whisper, torch; \
 
51
  print('--- Downloading sentence-transformers (120 MB) ---'); \
52
  SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L3-v2', cache_folder='/home/user/app/models_cache'); \
53
  print('--- Downloading TinyLlama tokenizer ---'); \
54
  AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0', cache_dir='/home/user/app/models_cache', trust_remote_code=True); \
55
- print('--- Downloading TinyLlama model weights (2.2 GB, please wait) ---'); \
56
- AutoModelForCausalLM.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0', cache_dir='/home/user/app/models_cache', torch_dtype=torch.float32, trust_remote_code=True); \
57
  print('--- Downloading Whisper base (140 MB) ---'); \
58
  whisper.load_model('base', download_root='/home/user/app/models_cache/whisper'); \
59
  print('=== All models downloaded successfully ==='); \
 
46
  # - Whisper base (~140 MB)
47
  RUN python -c "\
48
  from sentence_transformers import SentenceTransformer; \
49
+ from transformers import AutoTokenizer; \
50
+ from huggingface_hub import snapshot_download; \
51
+ import whisper; \
52
  print('--- Downloading sentence-transformers (120 MB) ---'); \
53
  SentenceTransformer('sentence-transformers/paraphrase-MiniLM-L3-v2', cache_folder='/home/user/app/models_cache'); \
54
  print('--- Downloading TinyLlama tokenizer ---'); \
55
  AutoTokenizer.from_pretrained('TinyLlama/TinyLlama-1.1B-Chat-v1.0', cache_dir='/home/user/app/models_cache', trust_remote_code=True); \
56
+ print('--- Downloading TinyLlama model weights to disk only (2.2 GB, no RAM loading) ---'); \
57
+ snapshot_download(repo_id='TinyLlama/TinyLlama-1.1B-Chat-v1.0', cache_dir='/home/user/app/models_cache'); \
58
  print('--- Downloading Whisper base (140 MB) ---'); \
59
  whisper.load_model('base', download_root='/home/user/app/models_cache/whisper'); \
60
  print('=== All models downloaded successfully ==='); \