Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -12,24 +12,26 @@ os.environ["TRANSFORMERS_VERBOSITY"] = "error"
|
|
12 |
|
13 |
# AutoModelForCausalLM과 AutoTokenizer를 로드합니다.
|
14 |
# BitNet 모델은 trust_remote_code=True가 필요합니다.
|
15 |
-
#
|
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(
|
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 인터페이스를 실행할 수 없습니다.")
|