import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from deep_translator import GoogleTranslator

# 1. 모델 및 토크나이저 로드
model_name = "SAVSNET/PetBERT_ICD"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# 2. 라벨 목록 (ICD 0~19 → 한글 번역)
LABELS = {
    0: "귀 또는 유양돌기 질환",
    1: "정신, 행동 또는 신경발달 장애",
    2: "혈액 또는 조혈기관 질환",
    3: "순환기계 질환",
    4: "치과 질환",
    5: "발달 이상",
    6: "소화기계 질환",
    7: "내분비, 영양 또는 대사 질환",
    8: "면역계 질환",
    9: "특정 감염성 또는 기생충 질환",
    10: "피부 질환",
    11: "근골격계 또는 결합조직 질환",
    12: "신생물(종양)",
    13: "신경계 질환",
    14: "시각계 질환",
    15: "주산기 기원의 특정 상태",
    16: "임신, 출산 또는 산후기 상태",
    17: "호흡기계 질환",
    18: "외상, 중독 또는 외부 원인 결과",
    19: "비뇨생식기계 질환"
}

# 3. 예측 함수 정의
def predict(text):
    try:
        # 번역 (한글 -> 영어)
        translated = GoogleTranslator(source='auto', target='en').translate(text)

        # 토큰화 및 모델 예측
        inputs = tokenizer(translated, return_tensors="pt", truncation=True, padding=True)
        with torch.no_grad():
            outputs = model(**inputs)
            logits = outputs.logits
            probs = torch.softmax(logits, dim=1).squeeze()

        # 예측 결과 상위 3개 추출
        topk = torch.topk(probs, 3)
        results = [
            f"{LABELS.get(int(idx), f'Label {idx}')} ({prob:.1%})"
            for idx, prob in zip(topk.indices, topk.values)
            if float(prob) > 0.1
        ]

        if results:
            summary_text = "예측된 질병:\n" + "\n".join(results)
        else:
            summary_text = "예측된 질병 없음 🫥"

        return summary_text

    except Exception as e:
        return f"오류 발생: {str(e)}"

# 4. Gradio UI 구성
demo = gr.Interface(
    fn=predict,
    inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 강아지가 자주 기침해"),
    outputs=gr.Textbox(label="예측 결과"),
    title="🐾 PetBERT ICD 수의사 예측기",
    description="반려동물의 증상 문장을 입력하면 AI가 질병 가능성을 예측해드립니다."
)

# 5. 앱 실행
demo.launch()