app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

# 모델명: 예시로 KR-FinBERT 감정분석 모델 사용 (금융 뉴스에 적합)
MODEL_NAME = "snunlp/KR-FinBERT"

# 모델/토크나이저 로딩 (최초 1회)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# 라벨 매핑 (모델에 따라 다를 수 있음)
id2label = {0: "neutral", 1: "positive", 2: "negative"}

def predict_sentiment(text):
    # 입력 전처리 및 토큰화
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=256)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = F.softmax(logits, dim=1).squeeze().tolist()
        pred_id = int(torch.argmax(logits, dim=1))
        label = id2label[pred_id]
        score = probs[pred_id]
    return {"label": label, "score": round(score, 4), "probs": {id2label[i]: round(p, 4) for i, p in enumerate(probs)}}

with gr.Blocks(title="국내 뉴스 감정분석 (KR-FinBERT)") as demo:
    gr.Markdown("# 국내 뉴스 감정분석 (KR-FinBERT)")
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(label="뉴스 본문", lines=12)
            submit_btn = gr.Button("Submit")
        with gr.Column():
            output_json = gr.JSON(label="분석 결과")
    submit_btn.click(predict_sentiment, input_text, output_json)

demo.launch()