neollm/utils/preprocess.py

import json
import re
from typing import Any, Callable


# dict2json --------------------------------
def dict2json(python_dict: dict[str, Any]) -> str:
    """
    Python dictをJSON文字列に変換する

    Args:
        python_dict (dict): 変換するPython dict

    Returns:
        str: 変換されたJSON文字列
    """
    # ensure_ascii: 日本語とかを出力するため
    json_string = json.dumps(python_dict, indent=2, ensure_ascii=False)
    return json_string


# optimize token --------------------------------
def optimize_token(text: str, funcs: list[Callable[[str], str]] | None = None) -> str:
    """
    テキストのトークンを最適化をする

    Args:
        text (str): 最適化するテキスト

    Returns:
        str: 最適化されたテキスト
    """
    funcs = funcs or [minimize_newline, zenkaku_to_hankaku, remove_trailing_spaces]
    for func in funcs:
        text = func(text)
    return text.strip()


def _replace_consecutive(text: str, pattern: str, replacing_text: str) -> str:
    """
    テキスト内の連続するパターンに対して、指定された置換テキストで置換する

    Args:
        text (str): テキスト
        pattern (str): 置換するパターン
        replacing_text (str): 置換テキスト

    Returns:
        str: 置換されたテキスト
    """
    p = re.compile(pattern)
    matches = [(m.start(), m.end()) for m in p.finditer(text)][::-1]

    text_replaced = list(text)

    for i_start, i_end in matches:
        text_replaced[i_start:i_end] = [replacing_text]
    return "".join(text_replaced)


def minimize_newline(text: str) -> str:
    """
    テキスト内の連続する改行を2以下にする

    Args:
        text (str): テキスト

    Returns:
        str: 改行を最小限にしたテキスト
    """
    return _replace_consecutive(text, pattern="\n{2,}", replacing_text="\n\n")


def zenkaku_to_hankaku(text: str) -> str:
    """
    テキスト内の全角文字を半角文字に変換する

    Args:
        text (str): テキスト

    Returns:
        str: 半角文字に変換されたテキスト
    """
    mapping_dict = {"　": " ", "：": ": ", "‎": " ", "．": "。", "，": "、", "￥": "¥"}
    hankaku_text = ""
    for char in text:
        # A-Za-z0-9!"#$%&'()*+,-./:;<=>?@[\]^_`{|}~
        if char in mapping_dict:
            hankaku_text += mapping_dict[char]
        elif 65281 <= ord(char) <= 65374:
            hankaku_text += chr(ord(char) - 65248)
        else:
            hankaku_text += char
    return hankaku_text


def remove_trailing_spaces(text: str) -> str:
    """
    テキスト内の各行の末尾のスペースを削除する

    Args:
        text (str): テキスト

    Returns:
        str: スペースを削除したテキスト
    """
    return "\n".join([line.rstrip() for line in text.split("\n")])