Update app.py

app.py

@@ -5,17 +5,22 @@ import pickle
 import gluonnlp as nlp
 import numpy as np
 import os
-import sys # 오류 시 서비스 종료를 위해 sys 모듈 임포트
+import sys
+import collections
+import logging # 로깅 모듈 임포트
 
 # transformers의 AutoTokenizer 및 BertModel 임포트
-from transformers import AutoTokenizer, BertModel # BertModel 임포트 추가
+from transformers import AutoTokenizer, BertModel
 from torch.utils.data import Dataset, DataLoader
-import logging # 로깅 모듈 임포트 유지
-from huggingface_hub import hf_hub_download # hf_hub_download 임포트 유지
-import collections # collections 모듈 임포트 유지
+from huggingface_hub import hf_hub_download
+
+# --- 로깅 설정 ---
+# INFO 레벨 이상의 로그를 출력하도록 설정합니다.
+# 실제 배포 환경에서는 로그 레벨을 WARNING이나 ERROR로 높여 불필요한 로그를 줄일 수 있습니다.
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
 
 # --- 1. BERTClassifier 모델 클래스 정의 ---
-# 이 클래스는 모델의 아키텍처를 정의합니다.
 class BERTClassifier(torch.nn.Module):
     def __init__(self,
                  bert,
@@ -49,11 +54,8 @@ class BERTClassifier(torch.nn.Module):
         return self.classifier(out)
 
 # --- 2. BERTDataset 클래스 정의 ---
-# 이 클래스는 데이터를 모델 입력 형식으로 변환합니다.
 class BERTDataset(Dataset):
     def __init__(self, dataset, sent_idx, label_idx, bert_tokenizer, vocab, max_len, pad, pair):
-        # nlp.data.BERTSentenceTransform은 토크나이저 함수를 받습니다.
-        # AutoTokenizer의 tokenize 메서드를 직접 전달합니다.
         transform = nlp.data.BERTSentenceTransform(
             bert_tokenizer, max_seq_length=max_len, vocab=vocab, pad=pad, pair=pair
         )
@@ -70,55 +72,45 @@ class BERTDataset(Dataset):
 app = FastAPI()
 device = torch.device("cpu") # Hugging Face Spaces의 무료 티어는 주로 CPU를 사용합니다.
 
-# ✅ category 로드 (GitHub 저장소 루트에 있어야 함)
+# ✅ category 로드
 try:
     with open("category.pkl", "rb") as f:
         category = pickle.load(f)
-    print("category.pkl 로드 성공.")
+    logger.info("category.pkl 로드 성공.")
 except FileNotFoundError:
-    print("Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
-    sys.exit(1) # 파일 없으면 서비스 시작하지 않음
+    logger.error("Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
+    sys.exit(1)
 
-# ✅ vocab 로드 (GitHub 저장소 루트에 있어야 함)
+# ✅ vocab 로드
 try:
     with open("vocab.pkl", "rb") as f:
         vocab = pickle.load(f)
-    print("vocab.pkl 로드 성공.")
+    logger.info("vocab.pkl 로드 성공.")
 except FileNotFoundError:
-    print("Error: vocab.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
-    sys.exit(1) # 파일 없으면 서비스 시작하지 않음
+    logger.error("Error: vocab.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
+    sys.exit(1)
 
-# ✅ 토크나이저 로드 (transformers.AutoTokenizer 사용)
+# ✅ 토크나이저 로드
 tokenizer = AutoTokenizer.from_pretrained('skt/kobert-base-v1')
-print("토크나이저 로드 성공.")
+logger.info("토크나이저 로드 성공.")
 
 # ✅ 모델 로드 (Hugging Face Hub에서 다운로드)
 try:
-    HF_MODEL_REPO_ID = "hiddenFront/TextClassifier" # 사용자님의 실제 Hugging Face 저장소 ID
-    HF_MODEL_FILENAME = "textClassifierModel.pt" # Hugging Face Hub에 업로드한 파일 이름과 일치해야 합니다.
+    HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
+    HF_MODEL_FILENAME = "textClassifierModel.pt"
     
     model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
-    print(f"모델 파일이 '{model_path}'에 성공적으로 다운로드되었습니다.")
+    logger.info(f"모델 파일이 '{model_path}'에 성공적으로 다운로드되었습니다.")
 
-    # --- 수정된 핵심 부분 ---
-    # 1. BertModel.from_pretrained를 사용하여 기본 BERT 모델을 로드합니다.
-    #    이렇게 하면 모델의 아키텍처와 사전 학습된 가중치가 로드됩니다.
     bert_base_model = BertModel.from_pretrained('skt/kobert-base-v1')
-
-    # 2. BERTClassifier 인스턴스를 생성합니다.
-    #    여기서 num_classes는 category 딕셔너리의 크기와 일치해야 합니다.
     model = BERTClassifier(
         bert_base_model,
         dr_rate=0.5, # 학습 시 사용된 dr_rate 값으로 변경하세요.
         num_classes=len(category)
     )
 
-    # 3. 다운로드된 파일에서 state_dict를 로드합니다.
-    #    이 파일은 사용자님의 경량화된 모델의 가중치만 포함하고 있습니다.
     loaded_state_dict = torch.load(model_path, map_location=device)
     
-    # 4. 로드된 state_dict의 키를 조정하고 모델에 적용합니다.
-    #    'module.' 접두사가 붙어있는 경우 제거하는 로직을 포함합니다.
     new_state_dict = collections.OrderedDict()
     for k, v in loaded_state_dict.items():
         name = k
@@ -126,19 +118,14 @@ try:
             name = name[7:]
         new_state_dict[name] = v
     
-    # strict=False를 사용하여 Missing key(s) 오류를 방지합니다.
-    # 이는 new_state_dict에 없는 키는 모델에서 기존 값(from_pretrained로 로드된)을 유지하고,
-    # 모델에 없는 키는 무시하도록 합니다.
     model.load_state_dict(new_state_dict, strict=False)
-    # --- 수정된 핵심 부분 끝 ---
-
-    model.to(device) # 모델을 디바이스로 이동
-    model.eval() # 추론 모드로 설정
-    print("모델 로드 성공.")
+    model.to(device)
+    model.eval()
+    logger.info("모델 로드 성공.")
 
 except Exception as e:
-    print(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
-    sys.exit(1) # 모델 로드 실패 시 서비스 시작하지 않음
+    logger.error(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
+    sys.exit(1)
 
 
 # ✅ 데이터셋 생성에 필요한 파라미터
@@ -149,26 +136,32 @@ batch_size = 32
 def predict(predict_sentence):
     data = [predict_sentence, '0']
     dataset_another = [data]
-    # num_workers는 배포 환경에서 0으로 설정 권장
-    # tokenizer.tokenize를 BERTDataset에 전달합니다.
     another_test = BERTDataset(dataset_another, 0, 1, tokenizer.tokenize, vocab, max_len, True, False)
     test_dataLoader = DataLoader(another_test, batch_size=batch_size, num_workers=0)
 
-    model.eval() # 예측 시 모델을 평가 모드로 설정
+    model.eval()
 
-    with torch.no_grad(): # 그라디언트 계산 비활성화
+    with torch.no_grad():
         for batch_id, (token_ids, valid_length, segment_ids, label) in enumerate(test_dataLoader):
             token_ids = token_ids.long().to(device)
             segment_ids = segment_ids.long().to(device)
             
             out = model(token_ids, valid_length, segment_ids)
             
-            logits = out
-            logits = logits.detach().cpu().numpy()
-            
-            predicted_category_index = np.argmax(logits)
-            predicted_category_name = list(category.keys())[predicted_category_index]
+            logits = out # 모델의 직접 출력은 로짓입니다.
+            probs = torch.nn.functional.softmax(logits, dim=1) # 확률 계산
+
+            predicted_category_index = torch.argmax(probs, dim=1).item() # 예측 인덱스
+            predicted_category_name = list(category.keys())[predicted_category_index] # 예측 카테고리 이름
             
+            # --- 예측 상세 로깅 ---
+            logger.info(f"Input Text: '{predict_sentence}'")
+            logger.info(f"Raw Logits: {logits.tolist()}")
+            logger.info(f"Probabilities: {probs.tolist()}")
+            logger.info(f"Predicted Index: {predicted_category_index}")
+            logger.info(f"Predicted Label: '{predicted_category_name}'")
+            # --- 예측 상세 로깅 끝 ---
+
             return predicted_category_name
 
 # ✅ 엔드포인트 정의