Create app.py

app.py

@@ -0,0 +1,91 @@
+from fastapi import FastAPI, Request
+from transformers import AutoTokenizer, BertForSequenceClassification, BertConfig
+from huggingface_hub import hf_hub_download
+import torch
+import numpy as np
+import pickle
+import sys
+import collections
+import os # os 모듈 임포트
+import psutil # 메모리 사용량 확인을 위해 psutil 임포트 (requirements.txt에 추가 필요)
+
+app = FastAPI()
+device = torch.device("cpu")
+
+# category.pkl 로드
+try:
+    with open("category.pkl", "rb") as f:
+        category = pickle.load(f)
+    print("category.pkl 로드 성공.")
+except FileNotFoundError:
+    print("Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
+    sys.exit(1)
+
+# 토크나이저 로드
+tokenizer = AutoTokenizer.from_pretrained("skt/kobert-base-v1")
+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")
+# --- 메모리 사용량 로깅 끝 ---
+
+try:
+    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
+    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")
+    # --- 메모리 사용량 로깅 끝 ---
+
+    model.eval()
+    print("모델 로드 성공.")
+except Exception as e:
+    print(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
+    sys.exit(1)
+
+@app.post("/predict")
+async def predict_api(request: Request):
+    data = await request.json()
+    text = data.get("text")
+    if not text:
+        return {"error": "No text provided", "classification": "null"}
+
+    encoded = tokenizer.encode_plus(
+        text, max_length=64, padding='max_length', truncation=True, return_tensors='pt'
+    )
+
+    with torch.no_grad():
+        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}