Update app.py

app.py

@@ -1,51 +1,42 @@
-import gradio as gr
-import torch
+from googletrans import Translator
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+import gradio as gr
 
-# ✅ 모델 로딩
-print("📦 모델 로딩 시작...")
+# 번역기 초기화
+translator = Translator()
+
+# PetBERT 모델 로딩
 tokenizer = AutoTokenizer.from_pretrained("SAVSNET/PetBERT_ICD")
 model = AutoModelForSequenceClassification.from_pretrained("SAVSNET/PetBERT_ICD")
-print("✅ 모델 로딩 완료!")
-
-# ✅ 예측 함수 정의
-def predict(text):
-    print("\n📥 입력 텍스트:", text)
 
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    print("🧪 Tokenized inputs:", inputs)
+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
-        print("📊 Logits:", logits)
-
-        probs = torch.sigmoid(logits).squeeze().tolist()
-        print("📈 확률값:", probs)
+        probs = torch.sigmoid(logits)
 
-    labels = model.config.id2label
-    print("🏷️ 라벨 목록:", labels)
+    threshold = 0.3  # 낮은 임계값
+    predicted = (probs > threshold).nonzero(as_tuple=True)[1]
 
-    result = {}
-    for i, p in enumerate(probs):
-        try:
-            label = labels[str(i)]  # ❗ str(i) 중요
-            if p > 0.5:
-                result[label] = f"{p:.1%}"
-        except Exception as e:
-            print(f"❌ 라벨 매핑 오류: {i} → {e}")
+    if len(predicted) == 0:
+        return {"예측된 질병 없음": "🐶"}
 
-    print("✅ 최종 예측 결과:", result)
-    return result or {"예측된 질병 없음": "🤷"}
+    result = {model.config.id2label[idx.item()]: f"{probs[0][idx].item():.2f}" for idx in predicted}
+    return result
 
-# ✅ Gradio 인터페이스 정의
+# Gradio 인터페이스
 demo = gr.Interface(
-    fn=predict,                    # 여기서 predict()가 호출됨
-    inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 기침을 자주 해요"),
-    outputs=gr.JSON(label="예측 질병 목록"),
-    title="🐾 PetBERT ICD 수의사 예측기",
-    description="반려동물의 증상 문장을 입력하면 AI가 질병 가능성을 예측합니다."
+    fn=predict_korean,
+    inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 강아지가 자주 기침해요."),
+    outputs="json",
+    title="🐾 PetBERT ICD 수의사 예측기 (한글 지원)",
+    description="한글로 반려동물 증상을 입력하면 영어로 번역한 뒤, AI가 질병을 예측합니다."
 )
 
-# ✅ 앱 실행
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()