Create model.py

model.py

@@ -0,0 +1,20 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import torch.nn.functional as F
+
+MODEL_NAME = "beomi/KcELECTRA-small"
+
+# 모델과 토크나이저 로드
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)  # fine-tuning 해야 정확해짐
+
+# 분류 함수
+def classify_text(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        probs = F.softmax(outputs.logits, dim=-1)
+        score = probs[0][1].item()  # 1번 라벨이 '욕설'이라고 가정
+
+    label = "욕설" if score > 0.5 else "정상"
+    return f"결과: {label} (욕설 확률: {score:.2%})"