train_tokenizer.py

# -*- coding: utf-8 -*-
import os
from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers

def train_tokenizer(texts, vocab_size, min_freq, output_dir, num_workers):
    """
    Εκπαίδευση ενός custom BPE tokenizer για το GPT-2.
    
    Args:
        texts (list): Λίστα με κείμενα για εκπαίδευση.
        vocab_size (int): Μέγεθος λεξιλογίου.
        min_freq (int): Ελάχιστη συχνότητα εμφάνισης για να συμπεριληφθεί token.
        output_dir (str): Φάκελος όπου θα αποθηκευτεί ο tokenizer.
        num_workers (int): Αριθμός εργαζομένων για παράλληλη επεξεργασία (σε αυτή την υλοποίηση δεν χρησιμοποιείται απευθείας).
    
    Returns:
        Tokenizer: Ο εκπαιδευμένος tokenizer.
    """
    # Αρχικοποίηση tokenizer με μοντέλο BPE και καθορισμό token για άγνωστα
    tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
    
    # Ορισμός του pre-tokenizer για byte-level επεξεργασία (όπως στο GPT-2)
    tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel()
    tokenizer.decoder = decoders.ByteLevel()
    
    # Ορισμός του trainer για BPE
    trainer = trainers.BpeTrainer(
        vocab_size=vocab_size,
        min_frequency=min_freq,
        special_tokens=["<s>", "<pad>", "</s>", "<unk>", "<mask>"]
    )
    
    # Εκπαίδευση του tokenizer με τα παρεχόμενα κείμενα
    tokenizer.train_from_iterator(texts, trainer=trainer, length=len(texts))
    
    # Δημιουργία φακέλου αποθήκευσης αν δεν υπάρχει
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Αποθήκευση του tokenizer σε αρχείο JSON
    tokenizer.save(os.path.join(output_dir, "tokenizer.json"))
    return tokenizer