build_markov.py

"""Build bigram AND trigram transition tables from the filtered word list."""

import json
from collections import defaultdict

import lang as _lang
from common import load_wordlist


def build_transitions() -> tuple[dict, dict]:
    """Build both bigram and trigram transition tables.

    Returns:
        (bigram_transitions, trigram_transitions)
        bigram: 2-char keys like "ab" -> {"c": 5, "d": 3}
        trigram: 3-char keys like "abc" -> {"d": 5, "e": 3}
    """
    words = load_wordlist()

    bigrams: defaultdict[str, defaultdict[str, int]] = defaultdict(lambda: defaultdict(int))
    trigrams: defaultdict[str, defaultdict[str, int]] = defaultdict(lambda: defaultdict(int))

    for word, freq in words:
        marked = f"^{word}$"
        for i in range(len(marked) - 2):
            bigram = marked[i : i + 2]
            next_char = marked[i + 2]
            bigrams[bigram][next_char] += freq

        for i in range(len(marked) - 3):
            trigram = marked[i : i + 3]
            next_char = marked[i + 3]
            trigrams[trigram][next_char] += freq

    def to_sorted_dict(d):
        return {
            key: dict(sorted(nexts.items(), key=lambda x: x[1], reverse=True))
            for key, nexts in sorted(d.items())
        }

    return to_sorted_dict(bigrams), to_sorted_dict(trigrams)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Build transition tables from word list")
    parser.add_argument(
        "--lang", required=True, choices=_lang.available_languages(), help="language code"
    )
    args = parser.parse_args()

    _lang.set_lang(args.lang)

    print(f"Building bigram and trigram transitions for {_lang.config()['name']}...")
    bigram_table, trigram_table = build_transitions()

    bigram_path = _lang.bigram_transitions_path()
    trigram_path = _lang.trigram_transitions_path()
    bigram_path.parent.mkdir(parents=True, exist_ok=True)

    with open(bigram_path, "w", encoding="utf-8") as f:
        json.dump(bigram_table, f, ensure_ascii=False, indent=2)

    with open(trigram_path, "w", encoding="utf-8") as f:
        json.dump(trigram_table, f, ensure_ascii=False, indent=2)

    print(f"Bigrams: {len(bigram_table)}")
    print(f"Bigram edges: {sum(len(v) for v in bigram_table.values())}")
    print(f"Trigrams: {len(trigram_table)}")
    print(f"Trigram edges: {sum(len(v) for v in trigram_table.values())}")
    print(f"Written to {bigram_path.parent}")

    print("\nSample trigrams:")
    shown = 0
    for key in trigram_table:
        if key.startswith("^") and len(key) == 3:
            top = list(trigram_table[key].items())[:5]
            print(f"  {key!r} -> {top}")
            shown += 1
            if shown >= 5:
                break