cutechicken commited on
Commit
b0b6eba
β€’
1 Parent(s): 44a7786

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +35 -4
app.py CHANGED
@@ -7,14 +7,37 @@ import json
7
  from datetime import datetime
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer
 
10
 
11
  # ν™˜κ²½ λ³€μˆ˜ μ„€μ •
12
  HF_TOKEN = os.getenv("HF_TOKEN")
13
  MODEL_ID = "CohereForAI/c4ai-command-r-plus-08-2024"
14
 
15
- # 영ꡬ μŠ€ν† λ¦¬μ§€ 경둜 μ„€μ •
16
- PERSISTENT_DIR = "/persistent/model_cache"
17
- os.makedirs(PERSISTENT_DIR, exist_ok=True)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
18
 
19
  # ν™˜κ²½ λ³€μˆ˜ μ„€μ •
20
  os.environ["TRANSFORMERS_CACHE"] = os.path.join(PERSISTENT_DIR, "transformers_cache")
@@ -27,7 +50,15 @@ for cache_dir in [
27
  os.environ["TORCH_HOME"],
28
  os.environ["HF_HOME"]
29
  ]:
30
- os.makedirs(cache_dir, exist_ok=True)
 
 
 
 
 
 
 
 
31
 
32
  class ModelManager:
33
  def __init__(self):
 
7
  from datetime import datetime
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer
10
+ import shutil
11
 
12
  # ν™˜κ²½ λ³€μˆ˜ μ„€μ •
13
  HF_TOKEN = os.getenv("HF_TOKEN")
14
  MODEL_ID = "CohereForAI/c4ai-command-r-plus-08-2024"
15
 
16
+ # μ‚¬μš©μž ν™ˆ 디렉토리 기반 경둜 μ„€μ •
17
+ USER_HOME = os.path.expanduser("~")
18
+ PERSISTENT_DIR = os.path.join(USER_HOME, "persistent_storage", "model_cache")
19
+
20
+ # μŠ€ν† λ¦¬μ§€ μƒνƒœ 확인
21
+ def check_storage():
22
+ try:
23
+ total, used, free = shutil.disk_usage(USER_HOME)
24
+ print(f"μŠ€ν† λ¦¬μ§€ μƒνƒœ:")
25
+ print(f"총 μš©λŸ‰: {total // (2**30)}GB")
26
+ print(f"μ‚¬μš© 쀑: {used // (2**30)}GB")
27
+ print(f"남은 μš©λŸ‰: {free // (2**30)}GB")
28
+ except Exception as e:
29
+ print(f"μŠ€ν† λ¦¬μ§€ 확인 μ‹€νŒ¨: {e}")
30
+
31
+ # 디렉토리 생성
32
+ try:
33
+ os.makedirs(PERSISTENT_DIR, exist_ok=True)
34
+ print(f"μΊμ‹œ 디렉토리 생성 μ™„λ£Œ: {PERSISTENT_DIR}")
35
+ except Exception as e:
36
+ print(f"μΊμ‹œ 디렉토리 생성 μ‹€νŒ¨: {e}")
37
+ # μž„μ‹œ λ””λ ‰ν† λ¦¬λ‘œ 폴백
38
+ PERSISTENT_DIR = os.path.join("/tmp", "model_cache")
39
+ os.makedirs(PERSISTENT_DIR, exist_ok=True)
40
+ print(f"μž„μ‹œ μΊμ‹œ 디렉토리 μ‚¬μš©: {PERSISTENT_DIR}")
41
 
42
  # ν™˜κ²½ λ³€μˆ˜ μ„€μ •
43
  os.environ["TRANSFORMERS_CACHE"] = os.path.join(PERSISTENT_DIR, "transformers_cache")
 
50
  os.environ["TORCH_HOME"],
51
  os.environ["HF_HOME"]
52
  ]:
53
+ try:
54
+ os.makedirs(cache_dir, exist_ok=True)
55
+ print(f"μƒμ„±λœ μΊμ‹œ 디렉토리: {cache_dir}")
56
+ except Exception as e:
57
+ print(f"μΊμ‹œ 디렉토리 생성 μ‹€νŒ¨: {cache_dir} - {e}")
58
+
59
+ # μ‹œμž‘ μ‹œ μŠ€ν† λ¦¬μ§€ μƒνƒœ 확인
60
+ check_storage()
61
+
62
 
63
  class ModelManager:
64
  def __init__(self):