app.py

from fastapi import FastAPI, Request
from pydantic import BaseModel
import torch
import pickle
import gluonnlp as nlp
import numpy as np
import os
import sys # 오류 시 서비스 종료를 위해 sys 모듈 임포트

# transformers의 AutoTokenizer만 사용합니다.
from transformers import AutoTokenizer # BertModel, BertForSequenceClassification 등은 이제 직접 필요 없습니다.
from torch.utils.data import Dataset, DataLoader
import logging # 로깅 모듈 임포트 유지
from huggingface_hub import hf_hub_download # hf_hub_download 임포트 유지
# collections 모듈은 더 이상 필요 없을 수 있지만, 혹시 몰라 유지합니다.
import collections 

# --- 1. FastAPI 앱 및 전역 변수 설정 ---
app = FastAPI()
device = torch.device("cpu") # Hugging Face Spaces의 무료 티어는 주로 CPU를 사용합니다.

# ✅ category 로드 (GitHub 저장소 루트에 있어야 함)
try:
    with open("category.pkl", "rb") as f:
        category = pickle.load(f)
    print("category.pkl 로드 성공.")
except FileNotFoundError:
    print("Error: category.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
    sys.exit(1) # 파일 없으면 서비스 시작하지 않음

# ✅ vocab 로드 (GitHub 저장소 루트에 있어야 함)
try:
    with open("vocab.pkl", "rb") as f:
        vocab = pickle.load(f)
    print("vocab.pkl 로드 성공.")
except FileNotFoundError:
    print("Error: vocab.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
    sys.exit(1) # 파일 없으면 서비스 시작하지 않음

# ✅ 토크나이저 로드 (transformers.AutoTokenizer 사용)
tokenizer = AutoTokenizer.from_pretrained('skt/kobert-base-v1')
print("토크나이저 로드 성공.")

# ✅ 모델 로드 (Hugging Face Hub에서 다운로드)
# textClassifierModel.pt 파일은 이미 경량화된 '완전한 모델 객체'라고 가정하고 직접 로드합니다.
try:
    HF_MODEL_REPO_ID = "hiddenFront/TextClassifier" # 사용자님의 실제 Hugging Face 저장소 ID
    HF_MODEL_FILENAME = "textClassifierModel.pt" # Hugging Face Hub에 업로드한 파일 이름과 일치해야 합니다.
    
    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
    print(f"모델 파일이 '{model_path}'에 성공적으로 다운로드되었습니다.")

    # --- 수정된 핵심 부분 ---
    # 경량화된 모델 객체를 직접 로드합니다.
    # 이 파일은 이미 PyTorch 모델 객체(양자화된 모델 포함)이므로 바로 로드하여 사용합니다.
    model = torch.load(model_path, map_location=device)
    # --- 수정된 핵심 부분 끝 ---

    model.eval() # 추론 모드로 설정
    print("모델 로드 성공.")

except Exception as e:
    print(f"Error: 모델 다운로드 또는 로드 중 오류 발생: {e}")
    sys.exit(1) # 모델 로드 실패 시 서비스 시작하지 않음


# --- 2. BERTDataset 클래스 정의 (dataset.py에서 옮겨옴) ---
# 이 클래스는 데이터를 모델 입력 형식으로 변환합니다.
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
        )
        self.sentences = [transform([i[sent_idx]]) for i in dataset]
        self.labels = [np.int32(i[label_idx]) for i in dataset]

    def __getitem__(self, i):
        return (self.sentences[i] + (self.labels[i],))

    def __len__(self):
        return len(self.labels)


# ✅ 데이터셋 생성에 필요한 파라미터
max_len = 64
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() # 예측 시 모델을 평가 모드로 설정

    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]
            
            return predicted_category_name

# ✅ 엔드포인트 정의
class InputText(BaseModel):
    text: str

@app.get("/")
def root():
    return {"message": "Text Classification API (KoBERT)"}

@app.post("/predict")
async def predict_route(item: InputText):
    result = predict(item.text)
    return {"text": item.text, "classification": result}