Update app.py

app.py

@@ -1,11 +1,11 @@
 from fastapi import FastAPI, Request
-from transformers import AutoTokenizer
+from transformers import BertForSequenceClassification, AutoTokenizer
 from huggingface_hub import hf_hub_download
 import torch
 import pickle
 import os
-import psutil
 import sys
+import psutil
 
 app = FastAPI()
 device = torch.device("cpu")
@@ -16,22 +16,27 @@ try:
         category = pickle.load(f)
     print("✅ category.pkl 로드 성공.")
 except FileNotFoundError:
-    print("❌ Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
+    print("❌ Error: category.pkl 파일을 찾을 수 없습니다.")
     sys.exit(1)
 
 # 토크나이저 로드
 tokenizer = AutoTokenizer.from_pretrained("skt/kobert-base-v1")
 print("✅ 토크나이저 로드 성공.")
 
+# 모델 구조 재정의
+num_labels = len(category)  # 분류할 클래스 수에 따라
+model = BertForSequenceClassification.from_pretrained("skt/kobert-base-v1", num_labels=num_labels)
+model.to(device)
+
 HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
 HF_MODEL_FILENAME = "textClassifierModel.pt"
 
-# 메모리 확인
+# 메모리 측정 전
 process = psutil.Process(os.getpid())
 mem_before = process.memory_info().rss / (1024 * 1024)
 print(f"📦 모델 다운로드 전 메모리 사용량: {mem_before:.2f} MB")
 
-# 모델 다운로드 및 로드
+# 모델 가중치 다운로드
 try:
     model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
     print(f"✅ 모델 파일 다운로드 성공: {model_path}")
@@ -39,14 +44,16 @@ try:
     mem_after_dl = process.memory_info().rss / (1024 * 1024)
     print(f"📦 모델 다운로드 후 메모리 사용량: {mem_after_dl:.2f} MB")
 
-    model = torch.load(model_path, map_location=device)  # 전체 모델 객체 로드
+    # state_dict 로드
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict)
     model.eval()
 
     mem_after_load = process.memory_info().rss / (1024 * 1024)
     print(f"📦 모델 로드 후 메모리 사용량: {mem_after_load:.2f} MB")
-    print("✅ 모델 로드 성공")
+    print("✅ 모델 로드 및 준비 완료.")
 except Exception as e:
-    print(f"❌ Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
+    print(f"❌ Error: 모델 로드 중 오류 발생: {e}")
     sys.exit(1)
 
 # 예측 API