core/tokenizer.py

"""
Core - Tokenizer Module

Compare différentes méthodes de tokenisation :
- NLTK word-level (punkt)
- SpaCy word-level
- HuggingFace subword WordPiece (BERT tokenizer)
- HuggingFace subword BPE (GPT-2 tokenizer)
"""

from typing import Any

import nltk
import spacy
from transformers import AutoTokenizer


# -- Téléchargement NLTK (une seule fois) --
try:
    nltk.data.find("tokenizers/punkt_tab")
except LookupError:
    nltk.download("punkt_tab", quiet=True)


# -- Chargement lazy des modèles --
_spacy_nlp = None
_hf_tokenizer = None
_gpt2_tokenizer = None


def _get_spacy_nlp():
    """Charge le modèle SpaCy en lazy loading."""
    global _spacy_nlp
    if _spacy_nlp is None:
        for model in ("fr_core_news_sm", "en_core_web_sm"):
            try:
                _spacy_nlp = spacy.load(model)
                return _spacy_nlp
            except OSError:
                continue
        # Auto-download si aucun modele trouve
        import subprocess
        subprocess.run(["python", "-m", "spacy", "download", "fr_core_news_sm"], check=True)
        _spacy_nlp = spacy.load("fr_core_news_sm")
    return _spacy_nlp


def _get_hf_tokenizer():
    """Charge le tokenizer BERT (WordPiece) en lazy loading."""
    global _hf_tokenizer
    if _hf_tokenizer is None:
        _hf_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
    return _hf_tokenizer


def _get_gpt2_tokenizer():
    """Charge le tokenizer GPT-2 (BPE) en lazy loading."""
    global _gpt2_tokenizer
    if _gpt2_tokenizer is None:
        _gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
    return _gpt2_tokenizer


# -- Fonctions de tokenisation --

def tokenize_nltk(text: str) -> dict[str, Any]:
    """Tokenise avec NLTK (word-level, punkt)."""
    tokens = nltk.word_tokenize(text)
    return {
        "method": "NLTK - Word-level (punkt)",
        "tokens": tokens,
        "count": len(tokens),
    }


def tokenize_spacy(text: str) -> dict[str, Any]:
    """Tokenise avec SpaCy (word-level)."""
    nlp = _get_spacy_nlp()
    doc = nlp(text)
    tokens = [token.text for token in doc]
    return {
        "method": "SpaCy - Word-level",
        "tokens": tokens,
        "count": len(tokens),
    }


def tokenize_huggingface(text: str) -> dict[str, Any]:
    """Tokenise avec BERT tokenizer (subword WordPiece)."""
    tokenizer = _get_hf_tokenizer()
    encoding = tokenizer(text, return_tensors=None)
    tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"])
    return {
        "method": "BERT - WordPiece",
        "tokens": tokens,
        "token_ids": encoding["input_ids"],
        "count": len(tokens),
    }


def tokenize_gpt2(text: str) -> dict[str, Any]:
    """Tokenise avec GPT-2 tokenizer (subword BPE byte-level)."""
    tokenizer = _get_gpt2_tokenizer()
    encoding = tokenizer(text, return_tensors=None)
    tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"])
    return {
        "method": "GPT-2 - BPE (byte-level)",
        "tokens": tokens,
        "token_ids": encoding["input_ids"],
        "count": len(tokens),
    }


def tokenize_all(text: str) -> list[dict[str, Any]]:
    """Exécute les 4 méthodes de tokenisation et retourne les résultats."""
    return [
        tokenize_nltk(text),
        tokenize_spacy(text),
        tokenize_huggingface(text),
        tokenize_gpt2(text),
    ]