Update app.py

app.py

@@ -23,10 +23,24 @@ except FileNotFoundError:
 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)
+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))
+
+    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 = CustomClassifier()
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.eval()
 
 HF_MODEL_REPO_ID = "hiddenFront/TextClassifier"
 HF_MODEL_FILENAME = "textClassifierModel.pt"