Edit app.py

app.py

@@ -1,34 +1,28 @@
 import gradio as gr
-# ... (토크나이저, 모델 로딩 코드) ...
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, TextClassificationPipeline
+import torch
 
-def predict_lyrics(lyrics_text):
-    print(f"API로부터 받은 가사: '{lyrics_text}'") # ★★★ 중요: 로그 확인용 ★★★
-    if not lyrics_text or not isinstance(lyrics_text, str) or lyrics_text.strip() == "":
-        print("입력 가사가 비어있거나 유효하지 않습니다.")
-        return "입력 가사가 비어있습니다. (신뢰도: 0%)" # 또는 다른 기본 응답
+MODEL_PATH = "./"
 
-    # ... (실제 모델 예측 로직) ...
-    # 예시: result = pipeline(lyrics_text)
-    #       label = result[0]['label']
-    #       score = result[0]['score']
-    #       return f"{label} (신뢰도: {score:.2f})"
-    return f"예측된 결과: {lyrics_text[::-1]}" # 임시 반환값 (실제 로직으로 대체)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_PATH)
+pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer, device=0 if torch.cuda.is_available() else -1)
 
-iface = gr.Interface(
-    fn=predict_lyrics,
-    inputs=gr.Textbox(lines=5, placeholder="가사를 입력하세요..."),
-    outputs=gr.Textbox(),
-    title="가사 장르 예측",
-    description="가사를 입력하면 장르를 예측해줍니다.",
-    # api_name="predict" # 명시적으로 API 이름을 지정할 수 있습니다. 기본값도 "predict"일 경우가 많습니다.
-)
+def predict(lyrics):
+    result = pipeline(lyrics)
+    # 예쁘게 결과만 추출
+    if isinstance(result, list) and len(result) > 0:
+        label = result[0].get("label", "Unknown")
+        score = result[0].get("score", 0)
+        return f"{label} ({score:.2f})"
+    return "No result"
 
-# 큐 설정을 통해 API를 열어둘 수 있습니다.
-# iface.queue(api_open=True) # 이 줄을 추가하거나, 아래 launch()에 queue()를 직접 호출
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(label="Lyrics"),
+    outputs=gr.Textbox(label="Predicted Genre"),
+    title="Lyrics Genre Predictor"
+)
 
 if __name__ == "__main__":
-    # iface.launch(share=True) # share=True는 외부 접속용이지만, hf spaces에서는 불필요
-    # queue()를 사용하여 API를 열어보세요.
-    iface.queue().launch() # 이렇게 하면 api_open=True가 기본값일 수 있습니다.
-                            # 또는 iface.launch(api_open=True) 같은 옵션은 Gradio 버전에 따라 다름
-                            # 명시적으로 iface.queue(api_open=True).launch()를 시도해보세요.
+    demo.launch()