Update app.py

app.py

@@ -1,13 +1,11 @@
 from fastapi import FastAPI, Request
-from transformers import AutoTokenizer, BertForSequenceClassification, BertConfig
+from transformers import AutoTokenizer
 from huggingface_hub import hf_hub_download
 import torch
-import numpy as np
 import pickle
+import os
+import psutil
 import sys
-import collections
-import os # os 모듈 임포트
-import psutil # 메모리 사용량 확인을 위해 psutil 임포트 (requirements.txt에 추가 필요)
 
 app = FastAPI()
 device = torch.device("cpu")
@@ -16,61 +14,42 @@ device = torch.device("cpu")
 try:
     with open("category.pkl", "rb") as f:
         category = pickle.load(f)
-    print("category.pkl 로드 성공.")
+    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("토크나이저 로드 성공.")
+print("✅ 토크나이저 로드 성공.")
 
 HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
 HF_MODEL_FILENAME = "textClassifierModel.pt"
 
-# --- 메모리 사용량 로깅 시작 ---
+# 메모리 확인
 process = psutil.Process(os.getpid())
-mem_before_model_download = process.memory_info().rss / (1024 * 1024) # MB 단위
-print(f"모델 다운로드 전 메모리 사용량: {mem_before_model_download:.2f} MB")
-# --- 메모리 사용량 로깅 끝 ---
+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}'에 성공적으로 다운로드되었습니다.")
+    print(f"✅ 모델 파일 다운로드 성공: {model_path}")
 
-    # --- 메모리 사용량 로깅 시작 ---
-    mem_after_model_download = process.memory_info().rss / (1024 * 1024) # MB 단위
-    print(f"모델 다운로드 후 메모리 사용량: {mem_after_model_download:.2f} MB")
-    # --- 메모리 사용량 로깅 끝 ---
-
-    # 1. 모델 아키텍처 정의 (가중치는 로드하지 않고 구조만 초기화)
-    config = BertConfig.from_pretrained("skt/kobert-base-v1", num_labels=len(category))
-    model = BertForSequenceClassification(config)
-
-    # 2. 다운로드된 파일에서 state_dict를 로드
-    loaded_state_dict = torch.load(model_path, map_location=device)
-
-    # 3. 로드된 state_dict를 정의된 모델에 적용
-    new_state_dict = collections.OrderedDict()
-    for k, v in loaded_state_dict.items():
-        name = k
-        if name.startswith('module.'):
-            name = name[7:]
-        new_state_dict[name] = v
-    
-    model.load_state_dict(new_state_dict)
-
-    # --- 메모리 사용량 로깅 시작 ---
-    mem_after_model_load = process.memory_info().rss / (1024 * 1024) # MB 단위
-    print(f"모델 로드 및 state_dict 적용 후 메모리 사용량: {mem_after_model_load:.2f} MB")
-    # --- 메모리 사용량 로깅 끝 ---
+    mem_after_dl = process.memory_info().rss / (1024 * 1024)
+    print(f"📦 모델 다운로드 후 메모리 사용량: {mem_after_dl:.2f} MB")
 
+    model = torch.load(model_path, map_location=device)  # 전체 모델 객체 로드
     model.eval()
-    print("모델 로드 성공.")
+
+    mem_after_load = process.memory_info().rss / (1024 * 1024)
+    print(f"📦 모델 로드 후 메모리 사용량: {mem_after_load:.2f} MB")
+    print("✅ 모델 로드 성공")
 except Exception as e:
-    print(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
+    print(f"❌ Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
     sys.exit(1)
 
+# 예측 API
 @app.post("/predict")
 async def predict_api(request: Request):
     data = await request.json()
@@ -86,6 +65,6 @@ async def predict_api(request: Request):
         outputs = model(**encoded)
         probs = torch.nn.functional.softmax(outputs.logits, dim=1)
         predicted = torch.argmax(probs, dim=1).item()
-    
+
     label = list(category.keys())[predicted]
     return {"text": text, "classification": label}