scripts/5_prepare_tf_token.py

import os
import sys
import json
from collections import Counter
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
from schemas import tf_token

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("フィールド別TF/トークン(tf_token.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.tf_token", "data/generated/tf_token.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)

    results: List[tf_token.Circle] = []

    for c in circles:
        circle_id = c.get("circleId")

        # 各フィールドのトークン化とTF
        field_objs: Dict[str, tf_token.TfOfField] = {}
        doc_tf_counter: Counter = Counter()
        doc_token_set: set = set()

        for field in target_fields:
            text = c.get(field) or ""
            toks = tokenize(str(text), tok, mode, target_pos_l1, ban_list, stopwords)
            tf = Counter(toks)
            field_objs[field] = tf_token.TfOfField(len=len(toks), tf=dict(tf))
            doc_tf_counter.update(tf)
            doc_token_set.update(tf.keys())

        # スキーマ Fields へ詰める（未定義フィールドは長さ0/空dictで埋める）
        def get_field(name: str) -> tf_token.TfOfField:
            return field_objs.get(name, tf_token.TfOfField(len=0, tf={}))

        fields_obj = tf_token.Fields(
            name=get_field("name"),
            circleName=get_field("circleName"),
            circleNameKana=get_field("circleNameKana"),
            description=get_field("description"),
            prSummary=get_field("prSummary"),
            prDetail=get_field("prDetail"),
        )

        circle_entry = tf_token.Circle(
            circleId=circle_id,
            fields=fields_obj,
            tf=dict(doc_tf_counter),
            # tokens はユニーク語彙の存在フラグ(1)とする
            tokens={t: 1 for t in sorted(doc_token_set)},
        )
        results.append(circle_entry)

    os.makedirs(os.path.dirname(output_path), exist_ok=True)
    json_dumps([r.model_dump() for r in results], output_path)
    log.info(f"tf_token.json を出力しました: {output_path}")


if __name__ == "__main__":
    main()