tokenizer.py

from pathlib import Path

import sentencepiece as spm


class NeedleTokenizer:
    pad_token_id = 0
    eos_token_id = 1
    bos_token_id = 2
    unk_token_id = 3
    tool_call_token_id = 4
    tools_token_id = 5

    def __init__(self, model_path):
        self.sp = spm.SentencePieceProcessor(model_file=str(model_path))

    @classmethod
    def from_pretrained(cls, model_dir):
        return cls(Path(model_dir) / "tokenizer.model")

    def encode(self, text):
        return self.sp.encode(text, out_type=int)

    def decode(self, tokens):
        return self.sp.decode(list(tokens))


def build_encoder_input(tokenizer, query, tools="[]", max_length=1024):
    query_tokens = tokenizer.encode(query)[: max_length - 2]
    remaining = max_length - len(query_tokens) - 1
    return query_tokens + [tokenizer.tools_token_id] + tokenizer.encode(tools)[:remaining]