Update app.py

app.py

@@ -27,7 +27,7 @@ def md_loading():
     tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-base')
     model = XLMRobertaForSequenceClassification.from_pretrained('xlm-roberta-base', num_labels=493)
     
-    model_checkpoint = 'base3_44_en.bin'
+    model_checkpoint = 'en_ko_4mix_proto.bin'
     project_path = './'
     output_model_file = os.path.join(project_path, model_checkpoint)
 
@@ -93,17 +93,16 @@ class TVT_Dataset(Dataset):
 
 
 # 텍스트 input 박스
-business = st.text_input('')
-# business_work = st.text_input('사업체 하는일')
-# work_department = st.text_input('근무부서')
-# work_position = st.text_input('직책')
-# what_do_i = st.text_input('내가 하는 일')
-
-business_work = ''
-work_department = ''
-work_position = ''
-what_do_i = ''
+business = st.text_input('사업체명')
+business_work = st.text_input('사업체 하는일')
+work_department = st.text_input('근무부서')
+work_position = st.text_input('직책')
+what_do_i = st.text_input('내가 하는 일')
 
+# business_work = ''
+# work_department = ''
+# work_position = ''
+# what_do_i = ''
 
 # data 준비
 
@@ -186,7 +185,9 @@ if st.button('확인'):
         # Move logits and labels to CPU
 #        logits = logits.detach().cpu().numpy()
 
-
+    pred_m = torch.nn.Softmax(dim=1)
+    pred_ = pred_m(logits)
+    # st.write(logits.size())
     # # 단독 예측 시
     # arg_idx = torch.argmax(logits, dim=1)
     # print('arg_idx:', arg_idx)
@@ -196,11 +197,15 @@ if st.button('확인'):
 
     # 상위 k번째까지 예측 시
     k = 10
-    topk_idx = torch.topk(logits.flatten(), k).indices    
+    topk_idx = torch.topk(pred_.flatten(), k).indices
+    topk_values = torch.topk(pred_.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 = topk_values.numpy()
+    
+    st.write(sum(torch.topk(pred_.flatten(), 493).values.numpy()))
     # print(num_ans, str_ans)
     # print(num_ans_topk)
 
@@ -224,16 +229,24 @@ if st.button('확인'):
     # print(str_ans, type(str_ans))
 
     str_ans_topk_list = []
+    percent_ans_topk_list = []
     for i in range(k):
         str_ans_topk_list.append(str_ans_topk[i].iloc[0])
+        percent_ans_topk_list.append(percent_ans_topk[i]*100)
 
     # print(str_ans_topk_list)
 
     ans_topk_df = pd.DataFrame({
         'NO': range(1, k+1),
         '세분류 코드': num_ans_topk,
-        '세분류 명칭': str_ans_topk_list
+        '세분류 명칭': str_ans_topk_list,
+        '확률': percent_ans_topk_list
     })
     ans_topk_df = ans_topk_df.set_index('NO')
 
-    st.dataframe(ans_topk_df)
+#    ans_topk_df.style.bar(subset='확률', align='left', color='blue')
+#    ans_topk_df['확률'].style.applymap(color='black', font_color='blue')
+
+#    st.dataframe(ans_topk_df)
+#    st.dataframe(ans_topk_df.style.bar(subset='확률', align='left', color='blue'))
+    st.write(ans_topk_df.style.bar(subset='확률', align='left', color='blue'))