Create utils.py

utils.py

@@ -0,0 +1,132 @@
+import numpy as np
+import torch
+from torch.nn import CrossEntropyLoss
+import itertools
+
+def is_hiragana_or_katakana(s):
+    for char in s:
+        if not ('\u3040' <= char <= '\u309F' or '\u30A0' <= char <= '\u30FF') or char == "ー":
+            return False
+    return True
+
+def add_dakuten_handakuten(query, string_type):
+    def convert_to_hiragana(s):
+        """与えられた文字列を平仮名に変換する"""
+        result = []
+        for char in s:
+            if 'ァ' <= char <= 'ヶ':  # 片仮名を平仮名に変換
+                result.append(chr(ord(char) - 96))
+            else:
+                result.append(char)
+        return ''.join(result)
+
+    def convert_to_katakana(s):
+        """与えられた文字列を片仮名に変換する"""
+        result = []
+        for char in s:
+            if 'ぁ' <= char <= 'ゖ':  # 平仮名を片仮名に変換
+                result.append(chr(ord(char) + 96))
+            else:
+                result.append(char)
+        return ''.join(result)
+
+    if string_type == "hiragana":
+        s = convert_to_hiragana(query)
+        dakuon_map = {
+            'か': 'が', 'き': 'ぎ', 'く': 'ぐ', 'け': 'げ', 'こ': 'ご',
+            'さ': 'ざ', 'し': 'じ', 'す': 'ず', 'せ': 'ぜ', 'そ': 'ぞ',
+            'た': 'だ', 'ち': 'ぢ', 'つ': 'づ', 'て': 'で', 'と': 'ど',
+            'は': 'ば', 'ひ': 'び', 'ふ': 'ぶ', 'へ': 'べ', 'ほ': 'ぼ'
+        }
+        handakuon_map = {
+            'は': 'ぱ', 'ひ': 'ぴ', 'ふ': 'ぷ', 'へ': 'ぺ', 'ほ': 'ぽ'
+        }
+    elif string_type == "katakana":
+        s = convert_to_katakana(query)
+        dakuon_map = {
+            'カ': 'ガ', 'キ': 'ギ', 'ク': 'グ', 'ケ': 'ゲ', 'コ': 'ゴ',
+            'サ': 'ザ', 'シ': 'ジ', 'ス': 'ズ', 'セ': 'ゼ', 'ソ': 'ゾ',
+            'タ': 'ダ', 'チ': 'ヂ', 'ツ': 'ヅ', 'テ': 'デ', 'ト': 'ド',
+            'ハ': 'バ', 'ヒ': 'ビ', 'フ': 'ブ', 'ヘ': 'ベ', 'ホ': 'ボ',
+            'ウ': 'ヴ'
+        }
+        handakuon_map = {
+            'ハ': 'パ', 'ヒ': 'ピ', 'フ': 'プ', 'ヘ': 'ペ', 'ホ': 'ポ'
+        }
+
+    # 文字ごとに元の文字と濁音・半濁音をリストにする
+    options = []
+    for char in s:
+        temp = [char]
+        if char in dakuon_map:
+            temp.append(dakuon_map[char])
+        if char in handakuon_map:
+            temp.append(handakuon_map[char])
+        options.append(temp)
+
+    # 全ての組み合わせを生成
+    candidates = list(itertools.product(*options))
+    return candidates
+
+def add_dashes(s):
+    if not s:
+        return ['']
+
+    # 再帰的に文字列の先頭以外の部分に「ー」を挿入するパターンを取得
+    substr_patterns = add_dashes(s[1:])
+
+    # 現在の文字を含めたパターンを生成
+    result = []
+    for pattern in substr_patterns:
+        result.append(s[0] + pattern)  # そのまま連結
+        result.append(s[0] + 'ー' + pattern)  # 「ー」を挿入して連結
+
+    return result
+
+def compute_losses(candidates, model, tokenizer):
+    inputs = tokenizer(candidates, return_tensors="pt", padding=True)
+    inputs["labels"] = inputs["input_ids"].masked_fill(inputs["input_ids"] == tokenizer.pad_token_id, -100)
+    inputs = inputs.to(model.device)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+        logits = outputs.logits
+        labels = inputs["labels"]
+
+        shift_logits = logits[..., :-1, :].contiguous()
+        shift_labels = labels[..., 1:].contiguous()
+        loss_fct = CrossEntropyLoss(reduction="none")
+
+        losses_flat = loss_fct(shift_logits.view(-1, model.config.vocab_size), shift_labels.view(-1))
+        losses_seq = losses_flat.view(shift_labels.shape)
+        mask_labels = shift_labels != tokenizer.pad_token_id
+        losses = torch.sum(losses_seq * mask_labels, -1) / mask_labels.sum(-1)
+
+    return losses
+
+def search_candidates(query, query_candidates, model, tokenizer, top_k=100):
+    old_query = query[:-1]
+    if old_query not in query_candidates:
+        old_candidates, _ = search_candidates(old_query, query_candidates, model=model, tokenizer=tokenizer, top_k=top_k)
+    else:
+        old_candidates, _ = query_candidates[old_query]
+        
+    string = query[-1]
+    candidates = []
+    for string_type in ["hiragana", "katakana"]:
+        candidates_ = add_dakuten_handakuten(string, string_type=string_type)
+        for candidate_ in candidates_:
+            candidates += add_dashes(candidate_)
+        
+    combinations = itertools.product(old_candidates, candidates)
+    new_candidates = [''.join(pair) for pair in combinations]
+    
+    losses = compute_losses(new_candidates, model=model, tokenizer=tokenizer)
+    sorted_items = torch.sort(losses)
+    sorted_candidates = np.array(new_candidates)[sorted_items.indices.cpu().numpy()]
+    topk_candidates = sorted_candidates[:top_k].tolist()
+    topk_losses = sorted_items.values[:top_k].cpu().tolist()
+    
+    query_candidates[query] = (topk_candidates, topk_losses)
+    return topk_candidates, topk_losses