from tokenizers import Tokenizer, decoders, models, normalizers, pre_tokenizers, trainers
from datasets import load_dataset
import re

# Προσαρμοσμένος Pre-tokenizer για κώδικα και πολυγλωσσικό κείμενο
code_regex = r"""(?x:
    //.*?$|/\*.*?\*/|  # Σχόλια
    "(?:\\.|[^\\"])*"|  # Strings
    '(?:\\.|[^\\'])*'|  # Chars
    \b(?:if|else|for|while|return|function)\b|  # Keywords
    [<>]=?|\+\+|--|&&|\|\||[-+*/%=&|^~!]=?|  # Operators
    \d+\.?\d*|\.\d+|  # Αριθμοί
    [{}[\](),.;:]|  # Σύμβολα
    \p{L}+|\p{N}+|  # Unicode letters/numbers
    \s+|  # Whitespace
    \S  # Άλλα
)"""

def train_tokenizer(iterator, vocab_size=32000, min_frequency=2):
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    
    # Προχωρημένο Normalization
    tokenizer.normalizer = normalizers.Sequence([
        normalizers.NFC(),
        normalizers.StripAccents()  # Προαιρετικό για τόνους
    ])
    
    # Προσαρμοσμένος Pre-tokenizer
    tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
        pre_tokenizers.Regex(code_regex),
        pre_tokenizers.ByteLevel(add_prefix_space=False)
    ])
    
    # Προχωρημένος Trainer
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_frequency,
        special_tokens=["<|endoftext|>", "<pad>", "<unk>", "<mask>"],
        continuing_subword_prefix="",
        show_progress=True
    )
    
    tokenizer.train_from_iterator(iterator, trainer=trainer)
    tokenizer.decoder = decoders.ByteLevel()
    return tokenizer