Update app.py

app.py

@@ -10,7 +10,7 @@ from sklearn.model_selection import train_test_split
 # 모델과 토크나이저 로드
 MODEL_NAME = "beomi/kcbert-base"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=2)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=3)  # 3개 클래스로 변경
 
 # 데이터셋 클래스 정의
 class CustomDataset(Dataset):
@@ -57,11 +57,12 @@ def train_model():
 "그녀는 친구들과 수다 떠는 것을 좋아한다.",여자
 "강력한 리더십으로 팀을 이끄는 모습이 인상적이었다.",남자
 "자신이 직접 만든 쿠키를 주변에 나누어주곤 한다.",여자
-"정민지",천사사
+"정민지",천사
 """
     
     data = pd.read_csv(io.StringIO(csv_data))
-    data['label'] = data['gender'].apply(lambda x: 0 if x == '남자' else 1)
+    # 3개 클래스로 라벨 변경: 남자=0, 여자=1, 천사=2
+    data['label'] = data['gender'].apply(lambda x: 0 if x == '남자' else (1 if x == '여자' else 2))
     train_data, _ = train_test_split(data, test_size=0.2, random_state=42)
     
     train_dataset = CustomDataset(train_data, tokenizer)
@@ -117,7 +118,10 @@ def predict_gender(text):
         prediction = torch.argmax(outputs.logits, dim=1).flatten().item()
         confidence = probabilities[0][prediction].item()
 
-    gender = "남자" if prediction == 0 else "여자"
+    # 3개 클래스 매핑: 0=남자, 1=여자, 2=천사
+    gender_map = {0: "남자", 1: "여자", 2: "천사"}
+    gender = gender_map[prediction]
+    
     return f"예측 성별: {gender} (신뢰도: {confidence:.2%})"
 
 # 앱 시작 시 모델 훈련
@@ -133,13 +137,15 @@ iface = gr.Interface(
         label="텍스트 입력"
     ),
     outputs=gr.Textbox(label="예측 결과"),
-    title="🤖 AI 성별 예측기",
-    description="입력된 텍스트를 바탕으로 성별을 예측합니다.",
+    title="🤖 AI 성별 예측기 (3분류)",
+    description="입력된 텍스트를 바탕으로 성별을 예측합니다. (남자/여자/천사)",
     examples=[
         ["그는 축구를 정말 좋아하고, 근육질의 몸매를 가졌다."],
         ["그녀는 긴 머리를 가졌고, 분홍색 원피스를 입었다."],
         ["짧은 머리에 정장을 입은 그는 회의에 참석했다."],
-        ["아름다운 목소리로 노래하는 그녀는 가수다."]
+        ["아름다운 목소리로 노래하는 그녀는 가수다."],
+        ["그들은 책 읽기를 좋아하고 조용한 성격이다."],
+        ["요리와 청소를 모두 잘하며 집안일을 도맡아 한다."]
     ],
     theme=gr.themes.Soft()
 )