kimhyunwoo commited on
Commit
ece9655
·
verified ·
1 Parent(s): 4bf6a80

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +7 -5
app.py CHANGED
@@ -12,24 +12,26 @@ os.environ["TRANSFORMERS_VERBOSITY"] = "error"
12
 
13
  # AutoModelForCausalLM과 AutoTokenizer를 로드합니다.
14
  # BitNet 모델은 trust_remote_code=True가 필요합니다.
15
- # bf16은 메모리 사용량을 줄이고 속도를 향상시킬 수 있습니다 (GPU 지원 시).
16
- # CPU만 사용하는 경우 torch_dtype을 생략하거나 torch.float32로 설정할 수 있습니다.
17
  try:
18
  print(f"모델 로딩 중: {model_id}...")
19
- tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
20
  # GPU가 사용 가능하면 bf16 사용
21
  if torch.cuda.is_available():
 
22
  model = AutoModelForCausalLM.from_pretrained(
23
  model_id,
24
  torch_dtype=torch.bfloat16,
25
  trust_remote_code=True
26
  ).to("cuda") # GPU로 모델 이동
 
27
  print("GPU를 사용하여 모델 로드 완료.")
28
  else:
 
29
  model = AutoModelForCausalLM.from_pretrained(
30
  model_id,
31
  trust_remote_code=True
32
  )
 
33
  print("CPU를 사용하여 모델 로드 완료. 성능이 느릴 수 있습니다.")
34
 
35
  except Exception as e:
@@ -86,7 +88,7 @@ if model is not None and tokenizer is not None:
86
  )
87
 
88
  # Gradio 앱 실행
89
- # share=True 하면 임시 공개 링크가 생성됩니다.
90
- interface.launch(share=False)
91
  else:
92
  print("모델 로드 실패로 인해 Gradio 인터페이스를 실행할 수 없습니다.")
 
12
 
13
  # AutoModelForCausalLM과 AutoTokenizer를 로드합니다.
14
  # BitNet 모델은 trust_remote_code=True가 필요합니다.
15
+ # GitHub 특정 브랜치에서 설치한 transformers를 사용합니다.
 
16
  try:
17
  print(f"모델 로딩 중: {model_id}...")
 
18
  # GPU가 사용 가능하면 bf16 사용
19
  if torch.cuda.is_available():
20
+ # torch_dtype을 명시적으로 설정하여 로드 오류 방지 시도
21
  model = AutoModelForCausalLM.from_pretrained(
22
  model_id,
23
  torch_dtype=torch.bfloat16,
24
  trust_remote_code=True
25
  ).to("cuda") # GPU로 모델 이동
26
+ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
27
  print("GPU를 사용하여 모델 로드 완료.")
28
  else:
29
+ # CPU 사용 시 torch_dtype 생략 또는 float32
30
  model = AutoModelForCausalLM.from_pretrained(
31
  model_id,
32
  trust_remote_code=True
33
  )
34
+ tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
35
  print("CPU를 사용하여 모델 로드 완료. 성능이 느릴 수 있습니다.")
36
 
37
  except Exception as e:
 
88
  )
89
 
90
  # Gradio 앱 실행
91
+ # Hugging Face Spaces에서는 share=True 자동으로 설정됩니다.
92
+ interface.launch()
93
  else:
94
  print("모델 로드 실패로 인해 Gradio 인터페이스를 실행할 수 없습니다.")