Add application file

app.py

@@ -1,17 +1,32 @@
-import os
+import gradio as gr
 from transformers import BertForSequenceClassification, BertTokenizer
+import torch
 
-def save_model_and_tokenizer(model, tokenizer, save_dir='my_kobert_model'):
-    if not os.path.exists(save_dir):
-        os.makedirs(save_dir)
-    model.save_pretrained(save_dir)
-    tokenizer.save_pretrained(save_dir)
-    print(f"모델과 토크나이저가 '{save_dir}' 폴더에 저장되었습니다.")
+# 모델과 토크나이저 불러오기 (본인 Hugging Face 모델 경로로 수정)
+MODEL_NAME = "young476/LyricToGenre0607"
+model = BertForSequenceClassification.from_pretrained(MODEL_NAME)
+tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
 
-if __name__ == '__main__':
-    # 이미 학습된 모델과 토크나이저 경로로 수정하세요
-    model_name_or_path = 'bert-base-uncased'  # 또는 직접 학습한 모델 경로
-    model = BertForSequenceClassification.from_pretrained(model_name_or_path)
-    tokenizer = BertTokenizer.from_pretrained(model_name_or_path)
+# 클래스 이름 리스트 (예시: 실제 장르명으로 수정)
+genre_labels = ["발라드", "댄스", "힙합", "록", "트로트", "R&B"]
 
-    save_model_and_tokenizer(model, tokenizer, save_dir='my_kobert_model')
+def predict_genre(lyrics):
+    inputs = tokenizer(lyrics, return_tensors="pt", truncation=True, max_length=256)
+    with torch.no_grad():
+        outputs = model(**inputs)
+        pred_id = outputs.logits.argmax(dim=1).item()
+        pred_label = genre_labels[pred_id]
+        probs = torch.softmax(outputs.logits, dim=1).squeeze().tolist()
+        prob_dict = {genre_labels[i]: float(probs[i]) for i in range(len(genre_labels))}
+    return pred_label, prob_dict
+
+demo = gr.Interface(
+    fn=predict_genre,
+    inputs=gr.Textbox(lines=8, label="가사 입력"),
+    outputs=[gr.Label(num_top_classes=1, label="예측 장르"), gr.Label(label="장르별 확률")],
+    title="가사 기반 장르 분류기",
+    description="한국 노래 가사를 입력하면 장르를 예측합니다."
+)
+
+if __name__ == "__main__":
+    demo.launch()