Update app.py

app.py

@@ -1,42 +1,56 @@
-from googletrans import Translator
+import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import gradio as gr
-
-# 번역기 초기화
-translator = Translator()
-
-# PetBERT 모델 로딩
-tokenizer = AutoTokenizer.from_pretrained("SAVSNET/PetBERT_ICD")
-model = AutoModelForSequenceClassification.from_pretrained("SAVSNET/PetBERT_ICD")
-
-def predict_korean(text_kr):
-    # 한글 → 영어 번역
-    translated = translator.translate(text_kr, src='ko', dest='en').text
-    print(f"🔁 번역된 문장: {translated}")
-
-    # 모델 입력 처리
-    inputs = tokenizer(translated, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        logits = model(**inputs).logits
-        probs = torch.sigmoid(logits)
-
-    threshold = 0.3  # 낮은 임계값
-    predicted = (probs > threshold).nonzero(as_tuple=True)[1]
-
-    if len(predicted) == 0:
-        return {"예측된 질병 없음": "🐶"}
-
-    result = {model.config.id2label[idx.item()]: f"{probs[0][idx].item():.2f}" for idx in predicted}
-    return result
-
-# Gradio 인터페이스
+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 코드 등은 예시, 필요시 HuggingFace card 참조)
+LABELS = {
+    0: "소화기 질환",
+    1: "호흡기 질환",
+    2: "피부 질환",
+    3: "정신 행동 이상",
+    4: "감염성 질환",
+    5: "기타 질병"
+}
+
+# 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)
+        result = {
+            LABELS.get(int(idx), f"Label {idx}") : f"{prob:.2%}"
+            for idx, prob in zip(topk.indices, topk.values)
+            if float(prob) > 0.1  # 10% 이상만 보여줌
+        }
+
+        return result if result else {"예측된 질병 없음": "🫥"}
+    except Exception as e:
+        return {"오류 발생": str(e)}
+
+# 4. Gradio UI 구성
 demo = gr.Interface(
-    fn=predict_korean,
-    inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 강아지가 자주 기침해요."),
-    outputs="json",
-    title="🐾 PetBERT ICD 수의사 예측기 (한글 지원)",
-    description="한글로 반려동물 증상을 입력하면 영어로 번역한 뒤, AI가 질병을 예측합니다."
+    fn=predict,
+    inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 강아지가 자주 기침해"),
+    outputs=gr.JSON(label="예측 질병 목록"),
+    title="🐾 PetBERT ICD 수의사 예측기",
+    description="반려동물의 증상 문장을 입력하면 AI가 질병 가능성을 예측해드립니다."
 )
 
+# 5. 앱 실행
 demo.launch()