[add] probability mode

utils/llm.py

@@ -19,6 +19,7 @@ def load_llm_from_pretrained(pretrained_model_path):
       device_map = {"": 0}
     else:
       device_map = {"": "cpu"}
+    print(device_map)
 
     # モデルの準備
     model_name = "line-corporation/japanese-large-lm-3.6b"
@@ -86,4 +87,29 @@ NGワード:
   )
   output = tokenizer.decode(output_ids.tolist()[0][input_ids.size(1):])
   
-  return output
+  return output
+
+def estimate_probability(model, tokenizer, input_text, ngword_with_label):
+  # プロンプトの準備
+  prompt = f"""指示:
+入力された文章が不適切（攻撃的、暴力的、差別的、性的、政治的）であるかどうか、「はい」か「いいえ」で回答してください。「はい」の場合は理由も添えてください。
+{input_text}
+NGワード:
+{", ".join(ngword_with_label)}
+応答:
+"""
+  input_ids = tokenizer.encode(prompt, add_special_tokens=False, return_tensors="pt")
+  with torch.no_grad():
+    output = model(input_ids)
+
+    next_token_logits = output.logits[0, -1, :]
+    next_token_probs = torch.softmax(next_token_logits, dim=-1)
+    sorted_ids = torch.argsort(next_token_probs, descending=True, dim=-1)
+  
+    for i in range(2):
+      prob_d[tokenizer.decode(sorted_ids[i])] = next_token_probs[sorted_ids[i]].item()
+    
+    if "はい" in prob_d:
+      return prob_d["はい"]
+    else:
+      return 0