app.py

from fastapi import FastAPI, Request
from transformers import BertForSequenceClassification, AutoTokenizer
from huggingface_hub import hf_hub_download
import torch
import pickle
import os
import sys
import psutil

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("✅ 토크나이저 로드 성공.")

# 모델 구조 재정의
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}")

    mem_after_dl = process.memory_info().rss / (1024 * 1024)
    print(f"📦 모델 다운로드 후 메모리 사용량: {mem_after_dl:.2f} MB")

    # 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("✅ 모델 로드 및 준비 완료.")
except Exception as e:
    print(f"❌ Error: 모델 로드 중 오류 발생: {e}")
    sys.exit(1)

# 예측 API
@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}