app.py

import html
import os
import re
from pathlib import Path

import gradio as gr
import torch
from peft import LoraConfig, PeftModel, TaskType
from transformers import AutoModelForSequenceClassification, AutoTokenizer

BASE_MODEL = "monologg/koelectra-small-v3-discriminator"
LORA_DIR = Path(__file__).parent / "lora_climate_misinfo"
MAX_LENGTH = 512
CONTACT_EMAIL = os.getenv("CONTACT_EMAIL", "").strip()
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 학습 데이터의 라벨 정의: 0=정상, 1=위험(가짜뉴스)
LABELS = {0: "정상", 1: "가짜뉴스 위험"}

LIMITATION_TEXT = (
    "이 모델은 AI 신약개발 관련 뉴스의 진위 판별을 돕는 보조 도구이며, 완벽한 사실검증기가 아닙니다. "
    "Model Card의 스트레스 테스트에서는 공신력 있는 기관명을 도용한 문장, 의약품과 무관한 공신력 기관명을 "
    "가져온 문장, 거짓 정보를 구체적 수치로 작성한 문장에서 오판 가능성이 확인되었습니다. "
    "특히 우회적·비꼬기 표현도 취약할 수 있으므로 최종 판단은 원문 출처와 논문·기관 자료를 함께 확인하세요."
)

CUSTOM_CSS = r"""
.gradio-container {max-width: 980px !important; margin: 0 auto !important;}
#hero {text-align:center; margin-bottom: 8px;}
.warning-box {border: 1px solid #e5b94f; background:#fff8df; border-radius:14px; padding:14px 16px; margin:8px 0 18px 0;}
.result-card {border:2px solid #222; border-radius:16px; padding:16px; background:white; min-height:132px;}
.result-head {font-size:14px; color:#555; margin-bottom:8px;}
.result-main {display:flex; align-items:center; gap:10px; font-size:24px; font-weight:700;}
.dot {width:18px; height:18px; border-radius:50%; display:inline-block; flex:0 0 auto;}
.dot-risk {background:#c95b4b;}
.dot-safe {background:#4e9f6a;}
.score-line {margin-top:10px; font-size:16px;}
.xai-box {border:2px solid #222; border-radius:20px; padding:14px; background:#fff;}
.xai-title {display:inline-block; border:2px solid #222; border-radius:20px; padding:2px 12px; font-weight:700; margin-bottom:12px;}
.token-wrap {display:flex; flex-wrap:wrap; gap:8px;}
.token-chip {border:1px solid #777; border-radius:12px; padding:6px 9px; background:#f7f7f7;}
.token-chip strong {font-weight:700;}
.small-note {font-size:13px; color:#666; margin-top:10px;}
.disclaimer {border:1px solid #999; border-radius:14px; padding:12px 14px; background:#f7f7f7; font-size:14px;}
.error-card {border:2px solid #a33; border-radius:14px; padding:14px; background:#fff3f3; color:#7b1f1f;}
"""


def preprocess_text(text: str) -> str:
    """흐름도의 '텍스트 전처리' 단계. 의미를 훼손하지 않는 범위에서 공백만 정리합니다."""
    text = (text or "").strip()
    text = re.sub(r"\s+", " ", text)
    return text


def build_peft_config() -> LoraConfig:
    """PEFT 0.9.0과 호환되는 LoRA 설정을 코드에서 명시적으로 생성합니다.

    업로드된 어댑터는 PEFT 0.20.0에서 저장되어 새 필드가 포함되어 있었기 때문에,
    구버전에서 'unexpected keyword'가 나지 않도록 필요한 설정만 사용합니다.
    """
    return LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["query", "value"],
        lora_dropout=0.1,
        bias="none",
        task_type=TaskType.SEQ_CLS,
        modules_to_save=["classifier", "score"],
        inference_mode=True,
        use_rslora=False,
        use_dora=False,
    )


def load_model_bundle():
    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, use_fast=True)
    base_model = AutoModelForSequenceClassification.from_pretrained(
        BASE_MODEL,
        num_labels=2,
    )
    peft_config = build_peft_config()
    model = PeftModel.from_pretrained(
        base_model,
        str(LORA_DIR),
        config=peft_config,
        is_trainable=False,
    )
    model.to(DEVICE)
    model.eval()
    return tokenizer, model


TOKENIZER = None
MODEL = None
MODEL_LOAD_ERROR = None
try:
    TOKENIZER, MODEL = load_model_bundle()
except Exception as exc:  # Space 자체가 죽지 않고 UI에서 원인을 확인할 수 있게 함
    MODEL_LOAD_ERROR = f"{type(exc).__name__}: {exc}"


def merge_wordpieces(tokens, scores):
    """WordPiece 토큰을 사람이 읽기 쉬운 단위로 묶어 상위 항목을 반환합니다."""
    merged = []
    for token, score in zip(tokens, scores):
        if token in {"[CLS]", "[SEP]", "[PAD]"}:
            continue
        if token.startswith("##") and merged:
            prev_token, prev_score = merged[-1]
            merged[-1] = (prev_token + token[2:], max(prev_score, float(score)))
        else:
            merged.append((token, float(score)))
    # 같은 표면형이 반복되면 가장 높은 점수만 유지
    best = {}
    for token, score in merged:
        token = token.strip()
        if not token or token in {"[UNK]"}:
            continue
        best[token] = max(score, best.get(token, -1.0))
    return sorted(best.items(), key=lambda x: x[1], reverse=True)


def extract_attention_guide(encoded, attentions, top_k=6):
    """마지막 층의 [CLS]→토큰 평균 attention을 간단한 참고 신호로 사용합니다.

    Attention은 인과적 설명이 아니므로 UI에도 그 한계를 명시합니다.
    """
    if not attentions:
        return []
    last = attentions[-1][0]  # [heads, seq, seq]
    cls_to_tokens = last[:, 0, :].mean(dim=0).detach().cpu().tolist()
    ids = encoded["input_ids"][0].detach().cpu().tolist()
    tokens = TOKENIZER.convert_ids_to_tokens(ids)
    merged = merge_wordpieces(tokens, cls_to_tokens)
    return merged[:top_k]


def result_html(pred_label: int, risk_prob: float, confidence: float):
    is_risk = pred_label == 1
    dot_class = "dot-risk" if is_risk else "dot-safe"
    title = "가짜뉴스 위험" if is_risk else "정상 가능성 높음"
    return f"""
    <div class='result-card'>
      <div class='result-head'>판정결과</div>
      <div class='result-main'><span class='dot {dot_class}'></span>{html.escape(title)} · {confidence:.3f}</div>
      <div class='score-line'><b>위험도 점수</b> {risk_prob * 100:.2f}%</div>
      <div class='small-note'>0=정상, 1=위험 라벨 기준의 모델 확률입니다. 확률값 자체가 사실의 증명은 아닙니다.</div>
    </div>
    """


def xai_html(guide):
    if not guide:
        chips = "<span class='token-chip'>추출된 토큰 없음</span>"
    else:
        max_score = max(score for _, score in guide) or 1.0
        chunks = []
        for token, score in guide:
            relative = score / max_score
            chunks.append(
                f"<span class='token-chip'><strong>{html.escape(token)}</strong> · attention {relative:.2f}</span>"
            )
        chips = "".join(chunks)
    return f"""
    <div class='xai-box'>
      <div class='xai-title'>Highlight Word · Attention Analysis</div>
      <div class='token-wrap'>{chips}</div>
      <div class='small-note'>위 토큰은 마지막 attention 층에서 상대적으로 크게 주목된 항목입니다. 인과적 근거나 사실검증 근거로 해석하면 안 됩니다.</div>
    </div>
    """


def analyze(text):
    cleaned = preprocess_text(text)
    if not cleaned:
        return (
            "<div class='error-card'>분석할 뉴스 문장을 입력하세요.</div>",
            xai_html([]),
            "입력 없음",
        )
    if MODEL_LOAD_ERROR:
        return (
            "<div class='error-card'><b>모델 로드 실패</b><br>Base model로 임의 판정하지 않습니다.<br>"
            + html.escape(MODEL_LOAD_ERROR)
            + "</div>",
            xai_html([]),
            "모델 로드 오류 — README의 런타임 점검 항목을 확인하세요.",
        )

    encoded = TOKENIZER(
        cleaned,
        return_tensors="pt",
        truncation=True,
        max_length=MAX_LENGTH,
    )
    encoded = {k: v.to(DEVICE) for k, v in encoded.items()}

    with torch.inference_mode():
        outputs = MODEL(
            **encoded,
            output_attentions=True,
            return_dict=True,
        )
        probs = torch.softmax(outputs.logits, dim=-1)[0].detach().cpu()

    pred = int(torch.argmax(probs).item())
    confidence = float(probs[pred].item())
    risk_prob = float(probs[1].item())
    guide = extract_attention_guide(encoded, outputs.attentions, top_k=6)

    status = (
        f"전처리 → 토크나이징 → KoELECTRA + LoRA 추론 → attention 기반 XAI 참고 → 결과 생성 완료. "
        f"판정: {LABELS[pred]}"
    )
    return result_html(pred, risk_prob, confidence), xai_html(guide), status


def appeal_message():
    if CONTACT_EMAIL:
        return f"이의 제기/오판 신고 연락처: {CONTACT_EMAIL}"
    return (
        "Model Card에는 개발자 이메일로 이의를 제기한다고 되어 있지만 실제 이메일 주소는 기재되어 있지 않습니다. "
        "배포 전 Hugging Face Space의 CONTACT_EMAIL 변수에 개발자 이메일을 등록하세요."
    )


with gr.Blocks(css=CUSTOM_CSS, title="AI 신약개발 뉴스 판별") as demo:
    gr.Markdown(
        "# AI 신약개발 뉴스 판별 모델\n"
        "AI 신약개발 관련 뉴스 문장을 입력하면 LoRA 분류 모델이 **정상/위험** 가능성을 표시하고, "
        "attention 기반 참고 토큰을 함께 보여줍니다.",
        elem_id="hero",
    )

    gr.HTML(f"<div class='warning-box'><b>⚠ Model Card 한계 고지</b><br>{html.escape(LIMITATION_TEXT)}</div>")

    with gr.Row():
        with gr.Column(scale=3):
            news_input = gr.Textbox(
                label="뉴스 텍스트",
                placeholder="AI 신약개발 관련 뉴스 문장 또는 기사 일부를 입력하세요.",
                lines=10,
            )
            analyze_btn = gr.Button("분석하기", variant="primary")
            gr.Examples(
                examples=[
                    ["연구진은 AI를 활용해 신약 후보 물질을 선별했으며, 실제 임상 효과는 추가 검증이 필요하다고 밝혔다."],
                    ["AI가 분석한 결과 이 치료제는 모든 암을 100% 치료하며 이미 세계 최고 기관이 효과를 보증했다."],
                ],
                inputs=news_input,
            )
        with gr.Column(scale=2):
            result = gr.HTML("<div class='result-card'><div class='result-head'>판정결과</div><div>분석 전</div></div>")
            status = gr.Textbox(label="처리 단계", value="대기 중", interactive=False)

    xai = gr.HTML(xai_html([]))

    gr.HTML(f"<div class='disclaimer'><b>최종 결과 및 한계 고지</b><br>{html.escape(LIMITATION_TEXT)}</div>")
    with gr.Row():
        appeal_btn = gr.Button("이의 제기 안내")
        appeal_output = gr.Textbox(label="이의 제기", interactive=False)

    with gr.Accordion("Model Card 요약", open=False):
        gr.Markdown(
            "- **모델명:** AI 신약 개발 분석 모델\n"
            "- **목적:** AI 신약개발 관련 뉴스가 진짜뉴스인지 가짜뉴스인지 구분하기 위한 보조 판별\n"
            "- **학습 데이터:** 동아사이언스 DB, 2024~2026년, 200건(정상 100 + 위험 100)\n"
            "- **Model Card 기재 성능:** Accuracy 100, F1 1.0, TP 20, TN 20, FP 0, FN 0\n"
            "- **가장 약한 유형:** 우회 공격\n"
            "- **금지된 사용:** 가짜뉴스 생성을 고도화하거나 모델을 우회하기 위한 용도"
        )

    analyze_btn.click(analyze, inputs=news_input, outputs=[result, xai, status])
    news_input.submit(analyze, inputs=news_input, outputs=[result, xai, status])
    appeal_btn.click(appeal_message, outputs=appeal_output)

if __name__ == "__main__":
    demo.queue().launch()