Update app.py

app.py

@@ -1,96 +1,72 @@
 from fastapi import FastAPI, Request
-from transformers import BertModel, BertForSequenceClassification, AutoTokenizer
-from huggingface_hub import hf_hub_download
+from pydantic import BaseModel
 import torch
 import pickle
+import gluonnlp as nlp
+import numpy as np
 import os
-import sys
-import psutil
+from kobert_tokenizer import KoBERTTokenizer
+from model import BERTClassifier
+from dataset import BERTDataset
+from transformers import BertModel
+import logging
 
 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)
+# ✅ category 로드
+with open("category.pkl", "rb") as f:
+    category = pickle.load(f)
 
-# 토크나이저 로드
-tokenizer = AutoTokenizer.from_pretrained("skt/kobert-base-v1")
-print("✅ 토크나이저 로드 성공.")
+# ✅ vocab 로드
+with open("vocab.pkl", "rb") as f:
+    vocab = pickle.load(f)
 
-class CustomClassifier(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        # 정의했던 구조 그대로 복원해야 함
-        self.bert = BertModel.from_pretrained("skt/kobert-base-v1")
-        self.classifier = torch.nn.Linear(768, len(category))
+# ✅ 토크나이저
+tokenizer = KoBERTTokenizer.from_pretrained('skt/kobert-base-v1')
 
-    def forward(self, input_ids, attention_mask=None, token_type_ids=None):
-        outputs = self.bert(input_ids=input_ids,
-                            attention_mask=attention_mask,
-                            token_type_ids=token_type_ids)
-        pooled_output = outputs[1]  # CLS 토큰
-        return self.classifier(pooled_output)
+# ✅ 모델 로드
+model = BERTClassifier(
+    BertModel.from_pretrained('skt/kobert-base-v1'),
+    dr_rate=0.5,
+    num_classes=len(category)
+)
+model.load_state_dict(torch.load("textClassifierModel.pt", map_location=device))
+model.to(device)
+model.eval()
 
-HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
-HF_MODEL_FILENAME = "textClassifierModel.pt"
+# ✅ 데이터셋 생성에 필요한 파라미터
+max_len = 64
+batch_size = 32
 
-# 메모리 측정 전
-process = psutil.Process(os.getpid())
-mem_before = process.memory_info().rss / (1024 * 1024)
-print(f"📦 모델 다운로드 전 메모리 사용량: {mem_before:.2f} MB")
+# ✅ 예측 함수
+def predict(predict_sentence):
+    data = [predict_sentence, '0']
+    dataset_another = [data]
+    another_test = BERTDataset(dataset_another, 0, 1, tokenizer, vocab, max_len, True, False)
+    test_dataLoader = torch.utils.data.DataLoader(another_test, batch_size=batch_size, num_workers=0)
 
-# 모델 로드
-try:
-    model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
-    print(f"✅ 모델 파일 다운로드 성공: {model_path}")
-
-    state_dict = torch.load(model_path, map_location=device)
-    model = BertForSequenceClassification.from_pretrained(
-        "skt/kobert-base-v1",
-        num_labels=len(category),
-        state_dict=state_dict,
-    )
-    model.to(device)
     model.eval()
-    print("✅ 모델 로드 및 준비 완료.")
-except Exception as e:
-    print(f"❌ Error: 모델 로드 중 오류 발생: {e}")
-    sys.exit(1)
+    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)
+        test_eval = []
+        for i in out:
+            logits = i.detach().cpu().numpy()
+            test_eval.append(list(category.keys())[np.argmax(logits)])
+        return test_eval[0]
 
+# ✅ 엔드포인트 정의
+class InputText(BaseModel):
+    text: str
 
 @app.get("/")
-def root(request: Request):
-    client_host = request.client.host
-    client_port = request.client.port
-    return {
-        "message": "Text Classification API is running!",
-        "client_ip": client_host,
-        "client_port": client_port
-    }
+def root():
+    return {"message": "Text Classification API (KoBERT)"}
 
-# 예측 API
 @app.post("/predict")
-async def predict_api(request: Request):
-    data = await request.json()
-    text = data.get("text")
-    print("request date", data);
-    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}
+async def predict_route(item: InputText):
+    result = predict(item.text)
+    return {"text": item.text, "classification": result}