tokenization_mentee.py

"""
Thin HuggingFace tokenizer wrapper around the custom BPE tokenizer.json.

Supports:
  AutoTokenizer.from_pretrained("MenteEAI/mentee-embed-v3", trust_remote_code=True)

Internally uses the `tokenizers` fast library — the same tokenizer.json
that was always shipped with the model.
"""
from __future__ import annotations

from transformers import PreTrainedTokenizerFast


class MenteeTokenizer(PreTrainedTokenizerFast):
    """
    Drop-in HuggingFace tokenizer for mentee-embed models.
    Wraps the BPE tokenizer.json trained alongside the model.

    Usage:
        from transformers import AutoTokenizer
        tok = AutoTokenizer.from_pretrained("MenteEAI/mentee-embed-v3", trust_remote_code=True)
        enc = tok(["hello world"], return_tensors="pt", padding=True, truncation=True)
    """

    # tell HF what the special tokens are named
    model_input_names = ["input_ids", "attention_mask"]

    def __init__(self, *args, **kwargs):
        # [PAD]=0, [UNK]=1 — these exist in the tokenizer.json vocab already
        kwargs.setdefault("pad_token", "[PAD]")
        kwargs.setdefault("unk_token", "[UNK]")
        super().__init__(*args, **kwargs)