scripts/4_prepare_bm25f_meta.py

import os
import sys
import json
import math
from collections import defaultdict
from typing import Dict, List

from sudachipy import dictionary, tokenizer

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from utils.logger import setup_logger
from utils.json import get_file_path_from_config, field_getter, json_dumps

log = setup_logger(__name__)


def load_configs():
    search = field_getter("config/search_model.json")
    files = field_getter("config/files.json")

    target_pos_l1: List[str] = search("target_pos_l1")
    target_fields: List[str] = search("target_fields")
    ban_list: List[str] = search("synonyms.banlist")

    stopwords_path = files("sudachi.stopwords")
    with open(stopwords_path, encoding="utf-8") as f:
        stopwords = set(json.load(f))

    sudachi_config_path = files("sudachi.sudachi_config")

    return target_pos_l1, target_fields, set(ban_list), stopwords, sudachi_config_path


def build_tokenizer(sudachi_config_path: str):
    tok = dictionary.Dictionary(config_path=sudachi_config_path).create()
    mode = tokenizer.Tokenizer.SplitMode.A
    return tok, mode


def tokenize(
    text: str, tok, mode, target_pos_l1: List[str], ban_list: set, stopwords: set
) -> List[str]:
    if not text:
        return []
    out: List[str] = []
    for m in tok.tokenize(text, mode):
        base = m.normalized_form().lower().strip()
        if not base:
            continue
        pos = m.part_of_speech()
        if pos[0] not in target_pos_l1:
            continue
        if base in stopwords or base in ban_list:
            continue
        out.append(base)
    return out


def main():
    log.info("BM25Fメタデータ(bm25_meta.json)を生成します")
    try:
        target_pos_l1, target_fields, ban_list, stopwords, sudachi_config_path = (
            load_configs()
        )
    except Exception as e:
        log.error(f"設定の読み込みに失敗しました: {e}")
        sys.exit(1)

    circles_path = get_file_path_from_config(
        "circles.circles_json", "data/generated/circles.json"
    )
    output_path = get_file_path_from_config(
        "bm25.bm25_meta", "data/generated/bm25_meta.json"
    )

    try:
        with open(circles_path, encoding="utf-8") as f:
            circles = json.load(f)
    except Exception as e:
        log.error(f"circles.jsonの読み込みに失敗しました: {e}")
        sys.exit(1)

    tok, mode = build_tokenizer(sudachi_config_path)

    N = len(circles)
    df: Dict[str, int] = defaultdict(int)
    field_token_lens_sum: Dict[str, int] = {f: 0 for f in target_fields}

    log.info(f"ドキュメント数: {N}")

    for c in circles:
        seen_in_doc = set()
        for field in target_fields:
            text = c.get(field) or ""
            if not text:
                continue
            toks = tokenize(str(text), tok, mode, target_pos_l1, ban_list, stopwords)
            field_token_lens_sum[field] += len(toks)
            for t in set(toks):
                if t not in seen_in_doc:
                    df[t] += 1
                    seen_in_doc.add(t)

    # IDF 計算（BM25で一般的な +0.5 smoothing と +1 オフセット）
    idf: Dict[str, float] = {}
    for term, dfi in df.items():
        idf_val = max(0.0, ((N - dfi + 0.5) / (dfi + 0.5)))
        # 数値安定化のためlog1p
        idf[term] = math.log1p(idf_val)

    avg_len = {
        field: (field_token_lens_sum[field] / N if N > 0 else 0.0)
        for field in target_fields
    }

    meta = {
        "N": N,
        "avg_len": avg_len,
        "idf": idf,
    }

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    json_dumps(meta, output_path)
    log.info(f"bm25_meta.json を出力しました: {output_path}")


if __name__ == "__main__":
    main()