cutechicken commited on
Commit
44a7786
·
verified ·
1 Parent(s): e5e0be0

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +19 -19
app.py CHANGED
@@ -7,21 +7,31 @@ import json
7
  from datetime import datetime
8
  import torch
9
  from transformers import AutoModelForCausalLM, AutoTokenizer
10
- import tempfile
11
 
12
  # 환경 변수 설정
13
  HF_TOKEN = os.getenv("HF_TOKEN")
14
  MODEL_ID = "CohereForAI/c4ai-command-r-plus-08-2024"
15
 
16
- # 임시 디렉토리 설정
17
- TEMP_DIR = tempfile.mkdtemp()
18
- os.environ["TRANSFORMERS_CACHE"] = os.path.join(TEMP_DIR, "transformers_cache")
19
- os.environ["TORCH_HOME"] = os.path.join(TEMP_DIR, "torch_cache")
20
- os.environ["HF_HOME"] = os.path.join(TEMP_DIR, "huggingface")
 
 
 
 
 
 
 
 
 
 
 
21
 
22
  class ModelManager:
23
  def __init__(self):
24
- self.cache_dir = os.path.join(TEMP_DIR, "model_cache")
25
  os.makedirs(self.cache_dir, exist_ok=True)
26
  self.model = None
27
  self.tokenizer = None
@@ -29,6 +39,7 @@ class ModelManager:
29
 
30
  def setup_model(self):
31
  try:
 
32
  print("토크나이저 로딩 시작...")
33
  self.tokenizer = AutoTokenizer.from_pretrained(
34
  MODEL_ID,
@@ -51,18 +62,7 @@ class ModelManager:
51
  print("모델 로딩 완료")
52
  except Exception as e:
53
  print(f"모델 로딩 중 오류 발생: {e}")
54
- try:
55
- print("기본 설정으로 다시 시도...")
56
- self.model = AutoModelForCausalLM.from_pretrained(
57
- MODEL_ID,
58
- token=HF_TOKEN,
59
- torch_dtype=torch.float16,
60
- trust_remote_code=True,
61
- cache_dir=self.cache_dir
62
- )
63
- print("기본 설정으로 모델 로딩 완료")
64
- except Exception as e2:
65
- raise Exception(f"모델 로딩 완전히 실패: {e2}")
66
 
67
  class ChatHistory:
68
  def __init__(self):
 
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")
21
+ os.environ["TORCH_HOME"] = os.path.join(PERSISTENT_DIR, "torch_cache")
22
+ os.environ["HF_HOME"] = os.path.join(PERSISTENT_DIR, "huggingface")
23
+
24
+ # 각 캐시 디렉토리 생성
25
+ for cache_dir in [
26
+ os.environ["TRANSFORMERS_CACHE"],
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):
34
+ self.cache_dir = os.path.join(PERSISTENT_DIR, "model_files")
35
  os.makedirs(self.cache_dir, exist_ok=True)
36
  self.model = None
37
  self.tokenizer = None
 
39
 
40
  def setup_model(self):
41
  try:
42
+ print(f"캐시 디렉토리 경로: {self.cache_dir}")
43
  print("토크나이저 로딩 시작...")
44
  self.tokenizer = AutoTokenizer.from_pretrained(
45
  MODEL_ID,
 
62
  print("모델 로딩 완료")
63
  except Exception as e:
64
  print(f"모델 로딩 중 오류 발생: {e}")
65
+ raise Exception(f"모델 로딩 실패: {e}")
 
 
 
 
 
 
 
 
 
 
 
66
 
67
  class ChatHistory:
68
  def __init__(self):