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

# 모델 로딩 (HuggingFace의 havocy28/VetBERT 사용)
tokenizer = AutoTokenizer.from_pretrained("havocy28/VetBERT")
model = AutoModelForSequenceClassification.from_pretrained("havocy28/VetBERT")

def classify_symptom(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        logits = model(**inputs).logits
        pred = torch.argmax(logits, dim=1).item()

    # 예시 레이블 이름 (모델에 맞춰 수정 가능)
    labels = {
        0: "정상",
        1: "소화기 질환",
        2: "호흡기 질환",
        3: "피부 질환",
        4: "기타 이상",
    }
    return labels.get(pred, f"예측 라벨: {pred}")

demo = gr.Interface(
    fn=classify_symptom,
    inputs=gr.Textbox(lines=3, placeholder="반려동물 증상을 입력해주세요."),
    outputs="text",
    title="VetBERT AI 수의사",
    description="반려견/묘 증상 문장을 입력하면 AI가 의심 질병을 분류해줍니다."
)

demo.launch()