Update app.py

app.py

@@ -195,11 +195,15 @@ if st.button('확인'):
 
     # 상위 k번째까지 예측 시
     k = 10
-    topk_idx = torch.topk(logits.flatten(), k).indices    
+    topk_idx = torch.topk(logits.flatten(), k).indices
+    topk_values = torch.topk(logits.flatten(), k).values
+
 
     num_ans_topk = label_tbl[topk_idx]
     str_ans_topk = [loc_tbl['항목명'][loc_tbl['코드'] == k] for k in num_ans_topk]
-
+    percent_ans_topk = list(topk_values)
+    
+    print(percent_ans_topk)
     # print(num_ans, str_ans)
     # print(num_ans_topk)
 
@@ -231,7 +235,8 @@ if st.button('확인'):
     ans_topk_df = pd.DataFrame({
         'NO': range(1, k+1),
         '세분류 코드': num_ans_topk,
-        '세분류 명칭': str_ans_topk_list
+        '세분류 명칭': str_ans_topk_list,
+        '확률': percent_ans_topk
     })
     ans_topk_df = ans_topk_df.set_index('NO')