api

type AnnotatedEncoding struct {
	IDs               []int       // token IDs
	Spans             []TokenSpan // byte spans for each token (use originalText[span.Start:span.End] to extract)
	SpecialTokensMask []int
}

type Config struct {
	ConfigFile     string
	TokenizerClass string `json:"tokenizer_class"`

	ChatTemplate           string `json:"chat_template"`
	UseDefaultSystemPrompt bool   `json:"use_default_system_prompt"`

	ModelMaxLength float64        `json:"model_max_length"`
	MaxLength      float64        `json:"max_length"`
	SpModelKwargs  map[string]any `json:"sp_model_kwargs"`

	ClsToken  string `json:"cls_token"`
	UnkToken  string `json:"unk_token"`
	SepToken  string `json:"sep_token"`
	MaskToken string `json:"mask_token"`
	BosToken  string `json:"bos_token"`
	EosToken  string `json:"eos_token"`
	PadToken  string `json:"pad_token"`

	AddBosToken             bool                  `json:"add_bos_token"`
	AddEosToken             bool                  `json:"add_eos_token"`
	AddedTokensDecoder      map[int]TokensDecoder `json:"added_tokens_decoder"`
	AdditionalSpecialTokens []string              `json:"additional_special_tokens"`

	DoLowerCase                bool `json:"do_lower_case"`
	CleanUpTokenizationSpaces  bool `json:"clean_up_tokenization_spaces"`
	SpacesBetweenSpecialTokens bool `json:"spaces_between_special_tokens"`

	TokenizeChineseChars bool   `json:"tokenize_chinese_chars"`
	StripAccents         any    `json:"strip_accents"`
	NameOrPath           string `json:"name_or_path"`
	DoBasicTokenize      bool   `json:"do_basic_tokenize"`
	NeverSplit           any    `json:"never_split"`

	Stride             int    `json:"stride"`
	TruncationSide     string `json:"truncation_side"`
	TruncationStrategy string `json:"truncation_strategy"`
}

type EncodeOptions struct {

	// AddSpecialTokens option takes a boolean value, and enables post-processing (e.g., [CLS]/[SEP] for BERT).
	// This is enabled by default for tokenizers that support it, since this is
	// the expected HuggingFace tokenizer behavior.
	AddSpecialTokens bool

	// MaxLen option takes an int value. Set it to a value <= 0 to disable MaxLen.
	// Encoding will be truncated to this length.
	MaxLen int

	// IncludeSpans option takes a boolean, and indicates if EncodeWithAnnotations should include spans.
	IncludeSpans bool

	// IncludeSpecialTokensMask option takes a boolean value, and enables post-processing (e.g., [CLS]/[SEP] for BERT).
	IncludeSpecialTokensMask bool
}

type SpecialToken int

type TokenSpan struct {
	Start int // start byte position (inclusive)
	End   int // end byte position (exclusive)
}

type Tokenizer interface {
	// Encode converts text to token IDs with post-processing (e.g., [CLS]/[SEP]).
	// Equivalent to EncodeWithOptions(text, true).
	Encode(text string) []int

	// EncodeWithAnnotations returns the encoded text along with requested annotations.
	// Use With to select the desired annotation.
	EncodeWithAnnotations(text string) AnnotatedEncoding

	// With applies options to a tokenizer.
	// If an option is unsupported, it returns an ErrNotImplemented derived error.
	//
	// These options may change how the tokenizer behaves, and usually they are set only once.
	// It is up to the caller to coordinate if the tokenizer is being used with alternating
	// options — it’s not a common pattern.
	With(options EncodeOptions) error

	// Decode converts the tokens back to their string representations.
	Decode([]int) string

	// SpecialTokenID returns ID for given special token if registered, or an ErrNotImplemented derived error if not.
	//
	// This allows for a common representation of special tokens across different tokenizers.
	SpecialTokenID(token SpecialToken) (int, error)

	// Normalize returns the normalization used by the tokenizer (e.g.: BERT lower cases the string)
	// If a Tokenizer uses no normalization this simply returns its input.
	Normalize(string) string

	// VocabSize returns the total number of tokens in the vocabulary.
	VocabSize() int

	// Config returns the HuggingFace tokenizer configuration.
	// It is optional, and in case the tokenizer has been instantiated in some other fashion it may return nil.
	Config() *Config
}

type TokensDecoder struct {
	Content    string `json:"content"`
	Lstrip     bool   `json:"lstrip"`
	Normalized bool   `json:"normalized"`
	Rstrip     bool   `json:"rstrip"`
	SingleWord bool   `json:"single_word"`
	Special    bool   `json:"special"`
}