import sys, time, json
from pathlib import Path
from transformers import AutoTokenizer

tok = AutoTokenizer.from_pretrained("gpt2")
p = Path(r"C:\RiemannModel\data\pretrain_corpus.txt")
regions = {}
total_tokens = 0
total_docs = 0
t0 = time.time()
batch = []


def flush(batch):
    global total_tokens
    enc = tok(batch, add_special_tokens=False)["input_ids"]
    n = sum(len(e) + 1 for e in enc)
    total_tokens += n
    return n


cum = 0
boundaries = []
with p.open("r", encoding="utf-8") as fh:
    for i, line in enumerate(fh):
        text = line.replace("\u2029", "\n").strip()
        if not text:
            continue
        batch.append(text)
        total_docs += 1
        if len(batch) == 2000:
            flush(batch)
            batch = []
        if total_docs % 500000 == 0:
            print(
                f"{total_docs:,} docs, {total_tokens/1e9:.3f}B tokens, {time.time()-t0:.0f}s",
                flush=True,
            )
            boundaries.append((total_docs, total_tokens))
if batch:
    flush(batch)
out = {
    "docs": total_docs,
    "tokens": total_tokens,
    "tokens_per_doc": total_tokens / max(1, total_docs),
    "cumulative_every_500k_docs": boundaries,
}
Path(r"C:\RiemannModel\data\corpus_token_count.json").write_text(
    json.dumps(out, indent=2)
)
print(
    json.dumps(
        {k: v for k, v in out.items() if k != "cumulative_every_500k_docs"}, indent=2
    )
)