import os
import tempfile

from datasets import load_dataset
from tokenizers import Tokenizer
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer
from tokenizers.pre_tokenizers import ByteLevel
from tokenizers.decoders import ByteLevel as ByteLevelDecoder
from tokenizers.processors import TemplateProcessing



BASE_DIR = r"E:\Aoban Selenia\Lightning"

PARQUET_FILE = os.path.join(
    BASE_DIR,
    "train0001of0001.parquet"
)

TOKENIZER_FILE = os.path.join(
    BASE_DIR,
    "lightning_tokenizer.json"
)

print(f"Loading dataset from {PARQUET_FILE}...")

dataset = load_dataset(
    "parquet",
    data_files=PARQUET_FILE,
    split="train"
)

print(f"Loaded {len(dataset)} rows.")


temp_file = tempfile.NamedTemporaryFile(
    mode="w",
    encoding="utf-8",
    suffix=".txt",
    delete=False
)

temp_path = temp_file.name

print(f"Creating temporary tokenizer dataset...")

try:

    with temp_file:

        for row in dataset:

            text = row["summary"]

            if text and text.strip():

                temp_file.write(text)
                temp_file.write("\n")


    print(f"Temporary dataset created at:")
    print(temp_path)



    tokenizer = Tokenizer(
        BPE(
            unk_token="<UNK>"
        )
    )

    tokenizer.pre_tokenizer = ByteLevel(
        add_prefix_space=True
    )

    tokenizer.decoder = ByteLevelDecoder()


    trainer = BpeTrainer(
        vocab_size=50000,
        min_frequency=2,
        special_tokens=[
            "<PAD>",
            "<UNK>",
            "<|endoftext|>"
        ]
    )



    print("Training BPE tokenizer...")

    tokenizer.train(
        files=[temp_path],
        trainer=trainer
    )


    tokenizer.post_processor = TemplateProcessing(
        single="$A <|endoftext|>",
        special_tokens=[
            (
                "<|endoftext|>",
                tokenizer.token_to_id("<|endoftext|>")
            )
        ]
    )


    tokenizer.save(
        TOKENIZER_FILE
    )

    print()
    print("Tokenizer trained successfully.")
    print(f"Vocab Size: {tokenizer.get_vocab_size()}")
    print(f"Saved to: {TOKENIZER_FILE}")


finally:


    if os.path.exists(temp_path):

        os.remove(temp_path)

        print("Temporary tokenizer dataset deleted.")