[add] probability mode

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 
 from utils.find_ng_word import get_ng_wordlist, get_ng_wordlist_from_saved, search_ng_word
-from utils.llm import load_llm_from_pretrained, inference
+from utils.llm import load_llm_from_pretrained, inference, estimate_probability
 
 
 wordlist_1_path_s = "input/ng_wordlists/ng_wordlist_1_sexual.txt"
@@ -53,11 +53,38 @@ def detect_ng_word(input_text):
       
     return response
 
+def estimate_ng_probability(input_text, thresold=0.3):
+    response = []
+    rtn_s = search_ng_word(input_text, ng_wordlist_1_s, ng_wordlist_2_s)
+    rtn_o = search_ng_word(input_text, ng_wordlist_1_o, ng_wordlist_2_o)
+    rtn = rtn_s + rtn_o
+
+    if len(rtn) == 0:
+        response.append("NGワードは検知されませんでした  \n")
+    else:
+        response.append('以下のNGワードを検知しました  \n')
+        for rtn_i in rtn:
+            ng_word = str(rtn_i) + "  \n"
+            response.append(ng_word)
+
+    rtn_s = [ri + "(sexual)" for ri in rtn_s]
+    rtn_o = [ri + "(offensive)" for ri in rtn_o]
+    ngword_with_label = rtn_s + rtn_o
+    
+    prob = estimate_probability(model, tokenizer, input_text, ngword_with_label)
+    
+    if prob > threshold:
+        response.append(f"不適切な内容を検知しました(NG度：{100*prob:.0f}%)")
+    else:
+        response.append(f"不適切な内容は検知されませんでした(NG度：{100*prob:.0f}%)")
+      
+    return response
+
 
 # 会話履歴用リスト型変数
 message_history = []
 
-def chat(user_msg):
+def chat(user_msg, mode, threshold):
     """
     AIとの会話を実行後、全会話履歴を返す
     user_msg: 入力されたユーザのメッセージ
@@ -71,7 +98,11 @@ def chat(user_msg):
     })
 
     # AIの回答を履歴に追加
-    response = detect_ng_word(user_msg)
+    if mode == "理由を出力":
+        response = detect_ng_word(user_msg)
+    elif mode == "確率を出力":
+        response = estimate_ng_probability(user_msg, threshold)
+
     assistant_msg = " ".join(response)
     message_history.append({
         "role": "assistant",
@@ -86,7 +117,12 @@ with gr.Blocks() as demo:
     # チャットボットUI処理
     chatbot = gr.Chatbot()
     input = gr.Textbox(show_label=False, placeholder="チェックしたい文章を入力してください")
-    input.submit(fn=chat, inputs=input, outputs=chatbot) # メッセージ送信されたら、AIと会話してチャット欄に全会話内容を表示
+    mode = gr.Radio(["理由を出力", "確率を出力"], label="モードを選択")
+    if mode == "確率を出力":
+        threshold = gr.Slider(minimum=0, maximum=1, value=0.3, step=0.05, label="NG度の閾値")
+    else:
+        threshold = None
+    input.submit(fn=chat, inputs=[input, mode, threshold], outputs=chatbot) # メッセージ送信されたら、AIと会話してチャット欄に全会話内容を表示
     input.submit(fn=lambda: "", inputs=None, outputs=input) # （上記に加えて）入力欄をクリア
 
-demo.launch()
+demo.launch(share=True)