tokenization_dreammachine.py

"""
DreamMachine "Tokenizer" — maps raw string IDs to multi-hash bucket indices.

HuggingFace convention requires a tokenizer file. For recommendation models,
the "tokenization" step is the multi-hash ID mapping.

Usage::

    from transformers import AutoTokenizer
    tok = AutoTokenizer.from_pretrained(
            "fangjunwen/dreammachine", trust_remote_code=True)

    item_ids = tok(item_id="SKU12345678")
    # → {"input_ids": tensor([[h1, h2, h3]])}  — K hash bucket indices

    user_ids = tok(user_id="U9876543210")
    # → {"input_ids": tensor([[h1, h2, h3]])}

    # Batch tokenise
    batch = tok(item_ids=["SKU001", "SKU002", "SKU003"])
    # → {"input_ids": tensor([[h1,h2,h3],[h1,h2,h3],[h1,h2,h3]])}
"""

import torch
import numpy as np
from typing import Dict, List, Optional, Union
from transformers import PreTrainedTokenizer


class DreamMachineTokenizer(PreTrainedTokenizer):
    """
    Multi-Hash ID tokenizer for DreamMachine.

    Converts arbitrary string IDs (item SKUs, user IDs, category names,
    brand names) into K hash bucket indices, matching the
    ``MultiHashEmbedding`` lookup in the model.

    This is not a text tokenizer — there is no vocabulary file.
    The ``vocab_size`` property returns ``hash_bucket_size``.

    Args:
        hash_bucket_size (int): B — number of hash buckets. Default: 50000.
        num_hashes (int):       K — number of independent hash functions. Default: 3.
    """

    vocab_files_names      = {}        # no vocab files needed
    model_input_names      = ["input_ids"]
    tokenizer_class        = "DreamMachineTokenizer"

    def __init__(
        self,
        hash_bucket_size: int = 50000,
        num_hashes: int = 3,
        unk_token: str = "<unk>",
        pad_token: str = "<pad>",
        **kwargs,
    ):
        self.hash_bucket_size = hash_bucket_size
        self.num_hashes       = num_hashes
        super().__init__(
            unk_token=unk_token,
            pad_token=pad_token,
            hash_bucket_size=hash_bucket_size,
            num_hashes=num_hashes,
            **kwargs,
        )

    @property
    def vocab_size(self) -> int:
        return self.hash_bucket_size

    def get_vocab(self) -> Dict[str, int]:
        # No discrete vocabulary — return empty dict
        return {}

    def _tokenize(self, text: str) -> List[str]:
        # Not used for hash-based IDs, but required by PreTrainedTokenizer
        return [text]

    def _convert_token_to_id(self, token: str) -> int:
        return abs(hash(f"{token}_seed_0")) % self.hash_bucket_size

    def _convert_id_to_token(self, index: int) -> str:
        return str(index)

    def convert_tokens_to_string(self, tokens: List[str]) -> str:
        return " ".join(tokens)

    # ── Main API ──────────────────────────────────────────────────────

    def hash_id(self, original_id: str) -> List[int]:
        """
        Map a single string ID to K hash bucket indices.

        Args:
            original_id: Any string (SKU, user_id, brand, category name …).

        Returns:
            List of K non-negative integer bucket indices.
        """
        return [
            abs(hash(f"{original_id}_seed_{k}")) % self.hash_bucket_size
            for k in range(self.num_hashes)
        ]

    def __call__(
        self,
        item_id: Optional[Union[str, List[str]]] = None,
        user_id: Optional[Union[str, List[str]]] = None,
        item_ids: Optional[List[str]] = None,
        user_ids: Optional[List[str]] = None,
        return_tensors: Optional[str] = "pt",
        **kwargs,
    ) -> Dict:
        """
        Tokenise one or more IDs.

        Args:
            item_id / user_id: Single string ID.
            item_ids / user_ids: Batch of string IDs.
            return_tensors: ``"pt"`` | ``"np"`` | ``None``.

        Returns:
            Dict with ``"input_ids"`` of shape ``[1, K]`` or ``[N, K]``.

        Examples::

            tok(item_id="SKU001")
            # {"input_ids": tensor([[12345, 67890, 34567]])}

            tok(item_ids=["SKU001", "SKU002"])
            # {"input_ids": tensor([[...], [...]])}
        """
        ids_list = item_ids or user_ids
        single   = item_id or user_id

        if single is not None:
            buckets = [self.hash_id(single)]
        elif ids_list is not None:
            buckets = [self.hash_id(i) for i in ids_list]
        else:
            raise ValueError("Provide item_id, user_id, item_ids, or user_ids.")

        arr = np.array(buckets, dtype=np.int64)  # [N, K]

        if return_tensors == "pt":
            return {"input_ids": torch.from_numpy(arr)}
        if return_tensors == "np":
            return {"input_ids": arr}
        return {"input_ids": arr.tolist()}

    # ── HF serialisation ──────────────────────────────────────────────

    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None):
        # No vocab file needed
        return ()

    def get_config(self) -> Dict:
        return {
            "hash_bucket_size": self.hash_bucket_size,
            "num_hashes":       self.num_hashes,
        }