Update app.py

app.py

@@ -5,13 +5,16 @@ import pickle
 import gluonnlp as nlp
 import numpy as np
 import os
-from kobert_tokenizer import KoBERTTokenizer # kobert_tokenizer 임포트 유지
-from transformers import BertModel # BertModel 임포트 유지
-from torch.utils.data import Dataset, DataLoader # DataLoader 임포트 추가
+import sys # sys 모듈 임포트 추가 (NameError 해결)
+
+# KoBERTTokenizer 대신 transformers.AutoTokenizer 사용
+from transformers import BertModel, AutoTokenizer # AutoTokenizer 임포트 유지
+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,
@@ -36,9 +39,6 @@ class BERTClassifier(torch.nn.Module):
     def forward(self, token_ids, valid_length, segment_ids):
         attention_mask = self.gen_attention_mask(token_ids, valid_length)
 
-        # BertModel의 출력 구조에 따라 수정
-        # Hugging Face Transformers의 BertModel은 (last_hidden_state, pooler_output, ...) 반환
-        # pooler_output (CLS 토큰의 최종 은닉 상태를 통과한 결과) 사용
         _, 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:
@@ -48,9 +48,10 @@ class BERTClassifier(torch.nn.Module):
         return self.classifier(out)
 
 # --- 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
         )
@@ -85,38 +86,30 @@ except FileNotFoundError:
     print("Error: vocab.pkl 파일을 찾을 수 없습니다. 프로젝트 루트에 있는지 확인하세요.")
     sys.exit(1) # 파일 없으면 서비스 시작하지 않음
 
-# ✅ 토크나이저 로드 (kobert_tokenizer 사용)
-# Colab 코드에서 사용된 방식이므로 유지합니다.
-tokenizer = KoBERTTokenizer.from_pretrained('skt/kobert-base-v1')
+# ✅ 토크나이저 로드 (transformers.AutoTokenizer 사용)
+# KoBERTTokenizer 대신 AutoTokenizer를 사용하여 KoBERT 모델의 토크나이저를 로드합니다.
+# 이렇게 하면 XLNetTokenizer 경고 및 kobert_tokenizer 설치 문제를 피할 수 있습니다.
+tokenizer = AutoTokenizer.from_pretrained('skt/kobert-base-v1')
 print("토크나이저 로드 성공.")
 
 # ✅ 모델 로드
-# 모델 아키텍처를 정의하고, 저장된 state_dict를 로드합니다.
 # num_classes는 category 딕셔너리의 크기와 일치해야 합니다.
+bertmodel = BertModel.from_pretrained('skt/kobert-base-v1')
 model = BERTClassifier(
-    BertModel.from_pretrained('skt/kobert-base-v1'),
+    bertmodel,
     dr_rate=0.5, # 학습 시 사용된 dr_rate 값으로 변경하세요.
     num_classes=len(category)
 )
 
 # textClassifierModel.pt 파일 로드
-# 이 파일은 GitHub 저장소에 없어야 하며, Dockerfile에서 Hugging Face Hub에서 다운로드하도록 설정되어 있습니다.
 try:
-    # Dockerfile에서 모델을 다운로드하도록 설정했으므로, 여기서는 로컬 경로를 사용합니다.
-    # 만약 Dockerfile에서 hf_hub_download를 사용하지 않는다면, 여기에 hf_hub_download를 추가해야 합니다.
-    # 현재 Dockerfile은 git+https://github.com/SKTBrain/KOBERT#egg=kobert_tokenizer 로드만 포함하고,
-    # 모델 파일 다운로드는 포함하지 않습니다.
-    # 따라서, 모델 파일을 Hugging Face Hub에서 다운로드하는 로직을 다시 추가해야 합니다.
-    from huggingface_hub import hf_hub_download
     HF_MODEL_REPO_ID = "hiddenFront/TextClassifier" # 사용자님의 실제 Hugging Face 저장소 ID
     HF_MODEL_FILENAME = "textClassifierModel.pt"
     model_path = hf_hub_download(repo_id=HF_MODEL_REPO_ID, filename=HF_MODEL_FILENAME)
     print(f"모델 파일이 '{model_path}'에 성공적으로 다운로드되었습니다.")
 
-    # 모델의 state_dict를 로드합니다.
     loaded_state_dict = torch.load(model_path, map_location=device)
     
-    # state_dict 키 조정 (필요한 경우)
     new_state_dict = collections.OrderedDict()
     for k, v in loaded_state_dict.items():
         name = k
@@ -143,7 +136,7 @@ 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)
+    another_test = BERTDataset(dataset_another, 0, 1, tokenizer, vocab, max_len, True, False) # tokenizer 객체 직접 전달
     test_dataLoader = DataLoader(another_test, batch_size=batch_size, num_workers=0)
 
     model.eval() # 예측 시 모델을 평가 모드로 설정