from tokenizer import load_tokenizer
from os import listdir
from tqdm import tqdm
import subprocess

if __name__ == "__main__":
    print("Preparing corpus for training...")

    tokenizer = load_tokenizer("trained/tokenizer_saved.json")

    print("Save path: data/pretraining/corpus_tokenized.txt")

    with open(
        "data/pretraining/corpus_tokenized.txt",
        "w",
        encoding="utf-8",
    ) as corpus_f:
        for filename in listdir("data/rawtexts/"):
            with open(
                "data/rawtexts/" + filename,
                "r",
                encoding="utf-8",
            ) as f:
                line = f.readline()
                with tqdm(
                    desc=f"Preprocessing corpus file: {filename}",
                    unit="lines",
                    total=int(
                        subprocess.run(
                            ["wc", "-l", "data/rawtexts/" + str(filename)],
                            capture_output=True,
                            text=True,
                        ).stdout.split()[0]
                    ),
                ) as prog:
                    while line != "":
                        encoding = tokenizer.encode(line)

                        for token_id in encoding.ids:
                            corpus_f.write(str(token_id))
                            corpus_f.write(" ")

                        line = f.readline()
                        prog.update(1)
    print("Done!")