bpe.go

package ztoken

import (
	"fmt"
	"math"
	"strings"
	"unicode/utf8"
)

// MergePair represents an adjacent token pair used in BPE merging.
//
// Stable.
type MergePair struct {
	Left  string
	Right string
}

// BPETokenizer implements the Tokenizer interface using byte-pair encoding.
// It loads vocabulary and merge rules from HuggingFace tokenizer.json format.
// When scores are set and merges are empty, it falls back to SentencePiece
// encoding using greedy leftmost-longest match with score-based tie-breaking,
// matching llama.cpp behavior.
//
// Stable.
type BPETokenizer struct {
	vocab        map[string]int
	reverseVocab map[int]string
	mergeRanks   map[MergePair]int
	special      SpecialTokens
	// byteEncoder maps each byte (0-255) to a printable Unicode character,
	// following the GPT-2 byte-level BPE convention.
	byteEncoder map[byte]rune
	// byteDecoder is the inverse of byteEncoder.
	byteDecoder map[rune]byte
	// preTokenize controls how text is split before BPE merging.
	// If true, byte-level pre-tokenization is used (GPT-2 style).
	byteLevelBPE bool
	// sentencePiece enables SentencePiece-style pre-tokenization where spaces
	// are replaced with ▁ (U+2581) and words are split at ▁ boundaries.
	sentencePiece bool
	// addLeadingSpace prepends ▁ to the first word during SentencePiece
	// pre-tokenization. This is true by default for SentencePiece models,
	// matching llama.cpp / SentencePiece behavior.
	addLeadingSpace bool
	// specialTokens maps special token strings to their IDs for exact matching
	// during encoding (e.g., "<start_of_turn>" -> 105).
	specialTokens map[string]int
	// normalizer is an optional text normalization function applied before tokenization.
	normalizer NormalizerFunc
	// scores holds SentencePiece unigram scores (negative log probabilities)
	// indexed by token ID. When scores are set and merges are empty, the
	// tokenizer uses greedy longest-match encoding instead of BPE merging.
	scores []float32
	// maxTokenLen caches the length (in bytes) of the longest token in vocab,
	// used to bound the search window in sentencePieceEncode.
	maxTokenLen int
}

// NewBPETokenizer creates a BPETokenizer from vocabulary, merge rules, and special tokens.
func NewBPETokenizer(vocab map[string]int, merges []MergePair, special SpecialTokens, byteLevelBPE bool) *BPETokenizer {
	reverseVocab := make(map[int]string, len(vocab))
	for k, v := range vocab {
		reverseVocab[v] = k
	}
	mergeRanks := make(map[MergePair]int, len(merges))
	for i, m := range merges {
		mergeRanks[m] = i
	}
	t := &BPETokenizer{
		vocab:        vocab,
		reverseVocab: reverseVocab,
		mergeRanks:   mergeRanks,
		special:      special,
		byteLevelBPE: byteLevelBPE,
	}
	if byteLevelBPE {
		t.byteEncoder, t.byteDecoder = buildByteEncoderDecoder()
	}
	return t
}

// Encode tokenizes text into a sequence of token IDs using BPE.
func (t *BPETokenizer) Encode(text string) ([]int, error) {
	if text == "" {
		return []int{}, nil
	}
	if t.normalizer != nil {
		text = t.normalizer(text)
	}

	// If special tokens are registered, split around them first.
	if len(t.specialTokens) > 0 {
		return t.encodeWithSpecials(text)
	}

	return t.encodeSegment(text, true)
}

// encodeSegment tokenizes a text segment that contains no special tokens.
// isFirstSegment indicates this is the first text segment (before any special
// tokens), which determines whether the addLeadingSpace field applies.
func (t *BPETokenizer) encodeSegment(text string, isFirstSegment bool) ([]int, error) {
	if text == "" {
		return nil, nil
	}
	var words []string
	if t.byteLevelBPE {
		words = t.byteLevelPreTokenize(text)
	} else if t.sentencePiece {
		words = t.sentencePiecePreTokenize(text, isFirstSegment && t.addLeadingSpace)
	} else {
		words = strings.Fields(text)
	}
	// When merges are empty but scores are available, use SentencePiece
	// unigram encoding (greedy longest-match) instead of BPE merging.
	useUnigram := len(t.mergeRanks) == 0 && len(t.scores) > 0

	var ids []int
	for _, word := range words {
		if useUnigram {
			ids = append(ids, t.sentencePieceEncode(word)...)
		} else {
			wordIDs, err := t.encodeWord(word)
			if err != nil {
				return nil, err
			}
			ids = append(ids, wordIDs...)
		}
	}
	return ids, nil
}

// encodeWithSpecials splits text around special token strings, encoding each
// special token as its single ID and encoding text between them with BPE.
// In SentencePiece mode, only the very first text segment (before any special
// token) gets the leading ▁ prefix. Text after special tokens does not.
func (t *BPETokenizer) encodeWithSpecials(text string) ([]int, error) {
	var ids []int
	isFirst := true
	for len(text) > 0 {
		// Find the earliest special token in the remaining text.
		bestIdx := -1
		bestLen := 0
		bestID := 0
		for tok, id := range t.specialTokens {
			idx := strings.Index(text, tok)
			if idx >= 0 && (bestIdx == -1 || idx < bestIdx || (idx == bestIdx && len(tok) > bestLen)) {
				bestIdx = idx
				bestLen = len(tok)
				bestID = id
			}
		}
		if bestIdx == -1 {
			// No more special tokens; encode the rest.
			segIDs, err := t.encodeSegment(text, isFirst)
			if err != nil {
				return nil, err
			}
			ids = append(ids, segIDs...)
			break
		}
		// Encode text before the special token.
		if bestIdx > 0 {
			segIDs, err := t.encodeSegment(text[:bestIdx], isFirst)
			if err != nil {
				return nil, err
			}
			ids = append(ids, segIDs...)
		}
		// Add the special token ID.
		ids = append(ids, bestID)
		isFirst = false
		text = text[bestIdx+bestLen:]
	}
	return ids, nil
}

// EncodeWithSpecialTokens wraps Encode and optionally prepends BOS / appends EOS.
func (t *BPETokenizer) EncodeWithSpecialTokens(text string, addBOS bool, addEOS bool) ([]int, error) {
	ids, err := t.Encode(text)
	if err != nil {
		return nil, err
	}
	if addBOS {
		ids = append([]int{t.special.BOS}, ids...)
	}
	if addEOS {
		ids = append(ids, t.special.EOS)
	}
	return ids, nil
}

// Decode converts token IDs back to text.
func (t *BPETokenizer) Decode(ids []int) (string, error) {
	var sb strings.Builder
	for _, id := range ids {
		tok, ok := t.reverseVocab[id]
		if !ok {
			return "", fmt.Errorf("unknown token ID: %d", id)
		}
		sb.WriteString(tok)
	}
	result := sb.String()
	if t.byteLevelBPE {
		decoded, err := t.decodeByteLevelBPE(result)
		if err != nil {
			return "", err
		}
		return decoded, nil
	}
	if t.sentencePiece {
		// Decode <0xNN> byte tokens back to actual bytes.
		result = decodeSentencePieceBytes(result)
		// Replace ▁ with space and trim leading space.
		result = strings.ReplaceAll(result, "\u2581", " ")
		result = strings.TrimPrefix(result, " ")
		return result, nil
	}
	return result, nil
}

// decodeSentencePieceBytes replaces <0xNN> hex byte tokens with the
// corresponding raw bytes. This reverses the byte fallback encoding
// used by SentencePiece for characters not in the vocabulary.
func decodeSentencePieceBytes(s string) string {
	var sb strings.Builder
	i := 0
	for i < len(s) {
		// Look for <0xNN> pattern: exactly 6 characters.
		if i+6 <= len(s) && s[i] == '<' && s[i+1] == '0' && s[i+2] == 'x' && s[i+5] == '>' {
			hi := unhex(s[i+3])
			lo := unhex(s[i+4])
			if hi >= 0 && lo >= 0 {
				sb.WriteByte(byte(hi<<4 | lo))
				i += 6
				continue
			}
		}
		sb.WriteByte(s[i])
		i++
	}
	return sb.String()
}

// unhex converts a hex digit character to its value, or -1 if invalid.
func unhex(c byte) int {
	switch {
	case c >= '0' && c <= '9':
		return int(c - '0')
	case c >= 'A' && c <= 'F':
		return int(c-'A') + 10
	case c >= 'a' && c <= 'f':
		return int(c-'a') + 10
	default:
		return -1
	}
}

// VocabSize returns the number of tokens in the vocabulary.
func (t *BPETokenizer) VocabSize() int {
	return len(t.vocab)
}

// GetToken returns the string for a given token ID.
func (t *BPETokenizer) GetToken(id int) (string, bool) {
	tok, ok := t.reverseVocab[id]
	return tok, ok
}

// GetID returns the ID for a given token string.
func (t *BPETokenizer) GetID(token string) (int, bool) {
	id, ok := t.vocab[token]
	return id, ok
}

// SpecialTokens returns the special token configuration.
func (t *BPETokenizer) SpecialTokens() SpecialTokens {
	return t.special
}

// SetSentencePiece enables SentencePiece-style pre-tokenization where spaces
// are replaced with ▁ (U+2581) and the text is split at ▁ boundaries.
func (t *BPETokenizer) SetSentencePiece(enabled bool) {
	t.sentencePiece = enabled
	t.addLeadingSpace = enabled
}

// SetAddLeadingSpace controls whether SentencePiece mode prepends ▁ to the
// first word. By default this is set to true when SetSentencePiece is called,
// matching llama.cpp / SentencePiece behavior. GGUF models may override this
// via the tokenizer.ggml.add_space_prefix metadata key.
func (t *BPETokenizer) SetAddLeadingSpace(enabled bool) {
	t.addLeadingSpace = enabled
}

// SetSpecialTokenStrings registers token strings that should be matched
// as single tokens during encoding instead of being split by BPE.
func (t *BPETokenizer) SetSpecialTokenStrings(tokens map[string]int) {
	t.specialTokens = tokens
}

// SetScores sets token scores for SentencePiece unigram encoding.
// When scores are set and merges are empty, the tokenizer uses
// score-based greedy encoding instead of BPE merge-based encoding.
// Scores are indexed by token ID (negative log probabilities).
func (t *BPETokenizer) SetScores(scores []float32) {
	t.scores = scores
	// Precompute max token length in bytes for search window bounding.
	t.maxTokenLen = 0
	for tok := range t.vocab {
		if len(tok) > t.maxTokenLen {
			t.maxTokenLen = len(tok)
		}
	}
}

// sentencePieceEncode tokenizes text using greedy leftmost-longest match.
//
// At each position, the longest vocabulary token is selected. Ties in length
// are broken by score (higher wins). This matches the llama.cpp SentencePiece
// tokenizer (llm_tokenizer_spm::tokenize) and produces the same output as
// HuggingFace for GGUF models. When no vocabulary token matches a byte,
// <0xNN> byte fallback tokens are used.
func (t *BPETokenizer) sentencePieceEncode(text string) []int {
	if text == "" {
		return nil
	}

	var ids []int
	pos := 0
	n := len(text)

	for pos < n {
		// Find the longest matching token at the current position.
		bestID := -1
		bestLen := 0
		bestScore := float32(math.Inf(-1))

		maxLen := t.maxTokenLen
		if maxLen > n-pos {
			maxLen = n - pos
		}

		for tokenLen := 1; tokenLen <= maxLen; tokenLen++ {
			candidate := text[pos : pos+tokenLen]
			if id, ok := t.vocab[candidate]; ok {
				score := t.tokenScore(id)
				// Prefer longest match; break ties by score.
				if tokenLen > bestLen || (tokenLen == bestLen && score > bestScore) {
					bestID = id
					bestLen = tokenLen
					bestScore = score
				}
			}
		}

		if bestID >= 0 {
			ids = append(ids, bestID)
			pos += bestLen
		} else {
			// Byte fallback: use <0xNN> for the current byte.
			byteToken := fmt.Sprintf("<0x%02X>", text[pos])
			if id, ok := t.vocab[byteToken]; ok {
				ids = append(ids, id)
			} else {
				ids = append(ids, t.special.UNK)
			}
			pos++
		}
	}

	return ids
}

// tokenScore returns the score for a token ID, or 0 if scores are not set
// or the ID is out of range.
func (t *BPETokenizer) tokenScore(id int) float32 {
	if id >= 0 && id < len(t.scores) {
		return t.scores[id]
	}
	return 0
}

// sentencePiecePreTokenize implements SentencePiece-style pre-tokenization.
// Text is split on whitespace boundaries. Words that follow a space get ▁
// (U+2581) prepended. Newlines are emitted as separate tokens.
// If addLeadingSpace is true, the very first word also gets ▁ prepended.
func (t *BPETokenizer) sentencePiecePreTokenize(text string, addLeadingSpace bool) []string {
	var words []string
	isFirstWord := true
	lines := strings.Split(text, "\n")
	for i, line := range lines {
		if i > 0 {
			words = append(words, "\n")
			isFirstWord = true // newline resets: next word has no ▁ prefix
		}
		if line == "" {
			continue
		}
		lineWords := strings.SplitAfter(line, " ")
		for _, w := range lineWords {
			w = strings.TrimRight(w, " ")
			if w == "" {
				continue
			}
			if isFirstWord && addLeadingSpace {
				words = append(words, "\u2581"+w)
			} else if !isFirstWord {
				words = append(words, "\u2581"+w)
			} else {
				words = append(words, w)
			}
			isFirstWord = false
		}
	}
	return words
}

// byteLevelPreTokenize converts text to byte-level BPE tokens.
// Each byte of the UTF-8 encoding is mapped to a printable Unicode character.
// Whitespace is preserved as part of tokens (prefixed to the following word).
func (t *BPETokenizer) byteLevelPreTokenize(text string) []string {
	// Split on whitespace boundaries, preserving the space as prefix of next word.
	var words []string
	var current strings.Builder
	for i, r := range text {
		if r == ' ' || r == '\t' || r == '\n' || r == '\r' {
			if current.Len() > 0 {
				words = append(words, current.String())
				current.Reset()
			}
			// Prefix space to next word token.
			if i == 0 || (i > 0 && (text[i-1] == ' ' || text[i-1] == '\t' || text[i-1] == '\n' || text[i-1] == '\r')) {
				// Leading/consecutive space becomes its own token.
				encoded := t.encodeBytesToChars([]byte{byte(r)})
				words = append(words, encoded)
			} else {
				current.WriteString(t.encodeBytesToChars([]byte{byte(r)}))
			}
		} else {
			current.WriteString(t.encodeBytesToChars([]byte(string(r))))
		}
	}
	if current.Len() > 0 {
		words = append(words, current.String())
	}
	return words
}

// encodeBytesToChars maps raw bytes to their BPE character representation.
func (t *BPETokenizer) encodeBytesToChars(b []byte) string {
	var sb strings.Builder
	for _, c := range b {
		sb.WriteRune(t.byteEncoder[c])
	}
	return sb.String()
}

// decodeByteLevelBPE reverses byte-level encoding back to UTF-8 text.
func (t *BPETokenizer) decodeByteLevelBPE(text string) (string, error) {
	var bytes []byte
	for _, r := range text {
		b, ok := t.byteDecoder[r]
		if !ok {
			return "", fmt.Errorf("unknown byte-level BPE character: %c (U+%04X)", r, r)
		}
		bytes = append(bytes, b)
	}
	return string(bytes), nil
}

// encodeWord applies BPE merging to a single pre-tokenized word.
func (t *BPETokenizer) encodeWord(word string) ([]int, error) {
	if word == "" {
		return nil, nil
	}

	// Split into individual characters as initial tokens.
	chars := []rune(word)
	symbols := make([]string, len(chars))
	for i, c := range chars {
		symbols[i] = string(c)
	}

	// Iteratively merge the highest-priority adjacent pair.
	for len(symbols) > 1 {
		bestRank := -1
		bestIdx := -1
		for i := 0; i < len(symbols)-1; i++ {
			pair := MergePair{Left: symbols[i], Right: symbols[i+1]}
			if rank, ok := t.mergeRanks[pair]; ok {
				if bestRank == -1 || rank < bestRank {
					bestRank = rank
					bestIdx = i
				}
			}
		}
		if bestIdx == -1 {
			break // No more merges possible.
		}
		// Merge the pair at bestIdx.
		merged := symbols[bestIdx] + symbols[bestIdx+1]
		newSymbols := make([]string, 0, len(symbols)-1)
		newSymbols = append(newSymbols, symbols[:bestIdx]...)
		newSymbols = append(newSymbols, merged)
		newSymbols = append(newSymbols, symbols[bestIdx+2:]...)
		symbols = newSymbols
	}

	// Look up token IDs for the final symbols.
	ids := make([]int, len(symbols))
	for i, sym := range symbols {
		id, ok := t.vocab[sym]
		if !ok {
			id = t.special.UNK
		}
		ids[i] = id
	}
	return ids, nil
}

// buildByteEncoderDecoder creates the GPT-2 byte-to-character mapping.
// Printable ASCII characters map to themselves. Other bytes map to
// Unicode characters starting at U+0100 (Latin Extended-B).
func buildByteEncoderDecoder() (map[byte]rune, map[rune]byte) {
	enc := make(map[byte]rune, 256)
	dec := make(map[rune]byte, 256)

	// Printable ASCII ranges that map to themselves:
	// '!' (33) to '~' (126), plus non-breaking characters.
	n := rune(256) // Next available Unicode codepoint for non-printable bytes.
	for i := range 256 {
		b := byte(i)
		if isPrintableGPT2Byte(b) {
			enc[b] = rune(b)
			dec[rune(b)] = b
		} else {
			enc[b] = n
			dec[n] = b
			n++
		}
	}
	return enc, dec
}

// isPrintableGPT2Byte returns true if the byte maps to itself in GPT-2 encoding.
func isPrintableGPT2Byte(b byte) bool {
	// '!' (33) through '~' (126)
	if b >= 33 && b <= 126 {
		return true
	}
	// Latin-1 supplement: 161-172, 174-255
	if b >= 161 && b <= 172 {
		return true
	}
	if b >= 174 {
		return true
	}
	return false
}

// decodeRune decodes the first UTF-8 rune from b and returns it with its byte length.
// If b is empty or invalid, it returns utf8.RuneError and 1 to ensure forward progress.
func decodeRune(b []byte) (rune, int) {
	r, size := utf8.DecodeRune(b)
	if size == 0 {
		return utf8.RuneError, 1
	}
	return r, size
}

// Statically assert BPETokenizer implements Tokenizer.
var _ Tokenizer = (*BPETokenizer)(nil)