Update app.py

app.py

@@ -8,30 +8,31 @@ model_name = "SAVSNET/PetBERT_ICD"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSequenceClassification.from_pretrained(model_name)
 
-# 2. 라벨 목록 (ICD 코드 등은 예시)
-LABELS = {
-    0: "소화기 질환",
-    1: "호흡기 질환",
-    2: "피부 질환",
-    3: "정신 행동 이상",
-    4: "감염성 질환",
-    5: "기타 질병"
-}
-
-# 3. 예측 함수 정의
+# 2. 모델의 id2label에서 라벨 목록 추출
+raw_labels = model.config.id2label  # e.g., {'0': 'Digestive', '1': 'Respiratory', ...}
+LABELS = {int(k): v for k, v in raw_labels.items()}
+
+# 3. 콘솔에 출력
+print("🔍 PetBERT ICD 라벨 목록:")
+for i in range(len(LABELS)):
+    print(f"{i}: {LABELS[i]}")
+
+# 4. 화면에 출력할 라벨 목록 텍스트 준비
+label_info = "\n".join([f"**{i}**: {LABELS[i]}" for i in range(len(LABELS))])
+
+# 5. 예측 함수 정의
 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()
+            probs = torch.softmax(outputs.logits, dim=1).squeeze()
 
-        # 예측 결과 상위 3개 추출
+        # 상위 3개 결과 추출
         topk = torch.topk(probs, 3)
         results = [
             f"{LABELS.get(int(idx), f'Label {idx}')} ({prob:.1%})"
@@ -40,23 +41,28 @@ def predict(text):
         ]
 
         if results:
-            summary_text = "예측된 질병:\n" + "\n".join(results)
+            return "예측된 질병:\n" + "\n".join(results)
         else:
-            summary_text = "예측된 질병 없음 🫥"
-
-        return summary_text
+            return "예측된 질병 없음 🫥"
 
     except Exception as e:
         return f"오류 발생: {str(e)}"
 
-# 4. Gradio UI 구성
+# 6. Gradio UI 구성
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Textbox(label="반려동물 증상 입력", placeholder="예: 강아지가 자주 기침해"),
     outputs=gr.Textbox(label="예측 결과"),
     title="🐾 PetBERT ICD 수의사 예측기",
-    description="반려동물의 증상 문장을 입력하면 AI가 질병 가능성을 예측해드립니다."
+    description="반려동물의 증상 문장을 입력하면 AI가 질병 가능성을 예측해드립니다.",
+    examples=["고양이가 밥을 안 먹고 자꾸 토해요", "강아지가 숨을 헐떡이고 기침을 해요"],
+    live=False
 )
 
-# 5. 앱 실행
-demo.launch()
+# 7. Launch + 라벨 목록 함께 출력
+with gr.Blocks() as app:
+    gr.Markdown("### 📋 PetBERT ICD 질병 라벨 목록")
+    gr.Markdown(label_info)
+    demo.render()
+
+app.launch()