Update app.py

app.py

@@ -5,68 +5,19 @@ import pickle
 import gluonnlp as nlp
 import numpy as np
 import os
-import sys # sys 모듈 임포트 추가 (NameError 해결)
+import sys # 오류 시 서비스 종료를 위해 sys 모듈 임포트
 
-# KoBERTTokenizer 대신 transformers.AutoTokenizer 사용
-from transformers import BertModel, AutoTokenizer # AutoTokenizer 임포트 유지
+# 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 임포트 추가
-import collections # collections 모듈 임포트 유지
-
-# --- 1. BERTClassifier 모델 클래스 정의 (model.py에서 옮겨옴) ---
-class BERTClassifier(torch.nn.Module):
-    def __init__(self,
-                 bert,
-                 hidden_size = 768,
-                 num_classes=5, # 분류할 클래스 수 (category 딕셔너리 크기와 일치)
-                 dr_rate=None,
-                 params=None):
-        super(BERTClassifier, self).__init__()
-        self.bert = bert
-        self.dr_rate = dr_rate
-
-        self.classifier = torch.nn.Linear(hidden_size , num_classes)
-        if dr_rate:
-            self.dropout = torch.nn.Dropout(p=dr_rate)
-
-    def gen_attention_mask(self, token_ids, valid_length):
-        attention_mask = torch.zeros_like(token_ids)
-        for i, v in enumerate(valid_length):
-            attention_mask[i][:v] = 1
-        return attention_mask.float()
-
-    def forward(self, token_ids, valid_length, segment_ids):
-        attention_mask = self.gen_attention_mask(token_ids, valid_length)
-
-        _, pooler = self.bert(input_ids=token_ids, token_type_ids=segment_ids.long(), attention_mask=attention_mask.float().to(token_ids.device), return_dict=False)
-
-        if self.dr_rate:
-            out = self.dropout(pooler)
-        else:
-            out = pooler
-        return self.classifier(out)
+from huggingface_hub import hf_hub_download # hf_hub_download 임포트 유지
+# collections 모듈은 더 이상 필요 없을 수 있지만, 혹시 몰라 유지합니다.
+import collections 
 
-# --- 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)
-
-# --- 3. FastAPI 앱 및 전역 변수 설정 ---
+# --- 1. FastAPI 앱 및 전역 변수 설정 ---
 app = FastAPI()
-device = torch.device("cpu") # Render의 무료 티어는 주로 CPU를 사용합니다.
+device = torch.device("cpu") # Hugging Face Spaces의 무료 티어는 주로 CPU를 사용합니다.
 
 # ✅ category 로드 (GitHub 저장소 루트에 있어야 함)
 try:
@@ -87,38 +38,24 @@ except FileNotFoundError:
     sys.exit(1) # 파일 없으면 서비스 시작하지 않음
 
 # ✅ 토크나이저 로드 (transformers.AutoTokenizer 사용)
-# KoBERTTokenizer 대신 AutoTokenizer를 사용하여 KoBERT 모델의 토크나이저를 로드합니다.
-# 이렇게 하면 XLNetTokenizer 경고 및 kobert_tokenizer 설치 문제를 피할 수 있습니다.
 tokenizer = AutoTokenizer.from_pretrained('skt/kobert-base-v1')
 print("토크나이저 로드 성공.")
 
-# ✅ 모델 로드
-# num_classes는 category 딕셔너리의 크기와 일치해야 합니다.
-bertmodel = BertModel.from_pretrained('skt/kobert-base-v1')
-model = BERTClassifier(
-    bertmodel,
-    dr_rate=0.5, # 학습 시 사용된 dr_rate 값으로 변경하세요.
-    num_classes=len(category)
-)
-
-# textClassifierModel.pt 파일 로드
+# ✅ 모델 로드 (Hugging Face Hub에서 다운로드)
+# textClassifierModel.pt 파일은 이미 경량화된 '완전한 모델 객체'라고 가정하고 직접 로드합니다.
 try:
     HF_MODEL_REPO_ID = "hiddenFront/TextClassifier" # 사용자님의 실제 Hugging Face 저장소 ID
-    HF_MODEL_FILENAME = "textClassifierModel.pt"
+    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}'에 성공적으로 다운로드되었습니다.")
 
-    loaded_state_dict = torch.load(model_path, map_location=device)
-    
-    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)
-    model.to(device) # 모델을 디바이스로 이동
+    # --- 수정된 핵심 부분 ---
+    # 경량화된 모델 객체를 직접 로드합니다.
+    # 이 파일은 이미 PyTorch 모델 객체(양자화된 모델 포함)이므로 바로 로드하여 사용합니다.
+    model = torch.load(model_path, map_location=device)
+    # --- 수정된 핵심 부분 끝 ---
+
     model.eval() # 추론 모드로 설정
     print("모델 로드 성공.")
 
@@ -127,6 +64,25 @@ except Exception as 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
@@ -136,7 +92,8 @@ def predict(predict_sentence):
     data = [predict_sentence, '0']
     dataset_another = [data]
     # num_workers는 배포 환경에서 0으로 설정 권장
-    another_test = BERTDataset(dataset_another, 0, 1, tokenizer, vocab, max_len, True, False) # tokenizer 객체 직접 전달
+    # 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() # 예측 시 모델을 평가 모드로 설정