indigotf/common.py

import os
import re
import json

DECOR_LINE = re.compile(r"^[\s=\-_~*#.]{4,}$")


def clean_text(text):
    lines = [ln for ln in text.splitlines() if not DECOR_LINE.match(ln)]
    text = "\n".join(lines)
    text = re.sub(r"\n{3,}", "\n\n", text)
    return text.strip() + "\n"


def collect_text_files(paths):
    files = []
    for p in paths:
        if os.path.isdir(p):
            for root, _, names in os.walk(p):
                files.extend(os.path.join(root, n) for n in sorted(names) if n.lower().endswith(".txt"))
        else:
            files.append(p)
    return sorted(files)


def read_clean(path):
    with open(path, encoding="utf-8") as f:
        return clean_text(f.read())


def save_meta(base_path, config, vocab, step, val_loss, backend, tokenizer=None):
    meta = {
        "config": config,
        "vocab": vocab,
        "step": step,
        "val_loss": val_loss,
        "backend": backend,
        "tokenizer": tokenizer or {"type": "char"},
    }
    with open(os.path.splitext(base_path)[0] + "_meta.json", "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False)


def load_meta(path):
    with open(os.path.splitext(path)[0] + "_meta.json", encoding="utf-8") as f:
        return json.load(f)


def build_tokenizer(tokenizer_info, vocab):
    if tokenizer_info.get("type") == "bpe":
        from indigotf.bpe import BPETokenizer

        return BPETokenizer.from_state(tokenizer_info)
    from indigotf.tokenizer import CharTokenizer

    return CharTokenizer(vocab)