mirror of
https://github.com/ollama/ollama.git
synced 2026-09-21 05:28:00 -05:00
tokenizer, ml: remove dead code
Several pieces outlived the code that used them. The root tokenizer package implemented the GGUF-side vocabularies for the Go engine and the safetensors-to-GGUF converter; nothing has imported it since the converter went. ml/backend.go held the Go engine's Backend, Context and Tensor interfaces, with fs.Config existing only to be returned from them, and a single CUDA template instance under ml/backend/ggml survived the engine removal along with the gitattributes entries for that tree and the CI change-filter globs for it and for the long-gone llama/llama.cpp. From the image generation engine, an integration test group that no test registers, its build tag, and the StepBar progress widget remained. DeviceInfo.IsBetter has no caller at all. All of it goes. Tidying the module file drops the regexp2 dependency and leaves protobuf as an indirect requirement. The llama3.2 tokenizer fixtures stay: the MLX runner's tokenizer uses them for its GGML parity test.
This commit is contained in:
Jesse Gross
@@ -7,22 +7,9 @@ llama/**/*.cuh linguist-vendored
|
||||
llama/**/*.m linguist-vendored
|
||||
llama/**/*.metal linguist-vendored
|
||||
|
||||
ml/backend/**/*.c linguist-vendored
|
||||
ml/backend/**/*.h linguist-vendored
|
||||
ml/backend/**/*.cpp linguist-vendored
|
||||
ml/backend/**/*.hpp linguist-vendored
|
||||
ml/backend/**/*.cu linguist-vendored
|
||||
ml/backend/**/*.cuh linguist-vendored
|
||||
ml/backend/**/*.m linguist-vendored
|
||||
ml/backend/**/*.metal linguist-vendored
|
||||
ml/backend/**/*.comp linguist-vendored
|
||||
ml/backend/**/*.glsl linguist-vendored
|
||||
ml/backend/**/CMakeLists.txt linguist-vendored
|
||||
|
||||
app/webview linguist-vendored
|
||||
|
||||
llama/build-info.cpp linguist-generated
|
||||
ml/backend/ggml/ggml/src/ggml-metal/ggml-metal-embed.s linguist-generated
|
||||
|
||||
* text=auto
|
||||
*.go text eol=lf
|
||||
|
||||
@@ -49,8 +49,6 @@ jobs:
|
||||
'LLAMA_CPP_VERSION' \
|
||||
'MLX_VERSION' \
|
||||
'MLX_C_VERSION' \
|
||||
'llama/llama.cpp/**/*' \
|
||||
'ml/backend/ggml/ggml/**/*' \
|
||||
'x/imagegen/mlx/**' \
|
||||
'x/imagegen/mlx/**/*' \
|
||||
'x/mlxrunner/xgrammar/native/**' \
|
||||
|
||||
@@ -1,20 +0,0 @@
|
||||
package fs
|
||||
|
||||
import "iter"
|
||||
|
||||
type Config interface {
|
||||
Architecture() string
|
||||
String(string, ...string) string
|
||||
Uint(string, ...uint32) uint32
|
||||
Float(string, ...float32) float32
|
||||
Bool(string, ...bool) bool
|
||||
|
||||
Strings(string, ...[]string) []string
|
||||
Ints(string, ...[]int32) []int32
|
||||
Floats(string, ...[]float32) []float32
|
||||
Bools(string, ...[]bool) []bool
|
||||
|
||||
Len() int
|
||||
Keys() iter.Seq[string]
|
||||
Value(key string) any
|
||||
}
|
||||
@@ -21,7 +21,6 @@ require (
|
||||
github.com/agnivade/levenshtein v1.1.1
|
||||
github.com/charmbracelet/bubbletea v1.3.10
|
||||
github.com/charmbracelet/lipgloss v1.1.0
|
||||
github.com/dlclark/regexp2 v1.11.5
|
||||
github.com/emirpasic/gods/v2 v2.0.0-alpha
|
||||
github.com/klauspost/compress v1.18.3
|
||||
github.com/mattn/go-runewidth v0.0.16
|
||||
@@ -87,6 +86,6 @@ require (
|
||||
golang.org/x/net v0.46.0 // indirect
|
||||
golang.org/x/term v0.36.0
|
||||
golang.org/x/text v0.30.0
|
||||
google.golang.org/protobuf v1.34.1
|
||||
google.golang.org/protobuf v1.34.1 // indirect
|
||||
gopkg.in/yaml.v3 v3.0.1
|
||||
)
|
||||
|
||||
@@ -39,8 +39,6 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
|
||||
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
|
||||
github.com/dgryski/trifles v0.0.0-20200323201526-dd97f9abfb48 h1:fRzb/w+pyskVMQ+UbP35JkH8yB7MYb4q/qhBarqZE6g=
|
||||
github.com/dgryski/trifles v0.0.0-20200323201526-dd97f9abfb48/go.mod h1:if7Fbed8SFyPtHLHbg49SI7NAdJiC5WIA09pe59rfAA=
|
||||
github.com/dlclark/regexp2 v1.11.5 h1:Q/sSnsKerHeCkc/jSTNq1oCm7KiVgUMZRDUoRu0JQZQ=
|
||||
github.com/dlclark/regexp2 v1.11.5/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
|
||||
github.com/emirpasic/gods/v2 v2.0.0-alpha h1:dwFlh8pBg1VMOXWGipNMRt8v96dKAIvBehtCt6OtunU=
|
||||
github.com/emirpasic/gods/v2 v2.0.0-alpha/go.mod h1:W0y4M2dtBB9U5z3YlghmpuUhiaZT2h6yoeE+C1sCp6A=
|
||||
github.com/erikgeiser/coninput v0.0.0-20211004153227-1c3628e74d0f h1:Y/CXytFA4m6baUTXGLOoWe4PQhGxaX0KpnayAqC48p4=
|
||||
|
||||
@@ -137,10 +137,6 @@ func TestQuantization(t *testing.T) {
|
||||
runIntegrationGroup(t, "quantization")
|
||||
}
|
||||
|
||||
func TestImageGeneration(t *testing.T) {
|
||||
runIntegrationGroup(t, "image-generation")
|
||||
}
|
||||
|
||||
func testName(s string) string {
|
||||
return strings.NewReplacer("/", "~", " ", "_").Replace(s)
|
||||
}
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
//go:build integration && !fast && !release && !library && !imagegen
|
||||
//go:build integration && !fast && !release && !library
|
||||
|
||||
package integration
|
||||
|
||||
|
||||
-413
@@ -1,413 +0,0 @@
|
||||
package ml
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"math"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/fs"
|
||||
)
|
||||
|
||||
type Backend interface {
|
||||
// Close frees all memory associated with this backend
|
||||
Close()
|
||||
|
||||
Load(ctx context.Context, progress func(float32)) error
|
||||
|
||||
Config() fs.Config
|
||||
Get(name string) Tensor
|
||||
NewContext() Context
|
||||
NewContextSize(size int) Context
|
||||
|
||||
// Enumerate the devices available for inference via this backend
|
||||
BackendDevices() []DeviceInfo
|
||||
}
|
||||
|
||||
// BackendCacheConfig should be implemented by backends that need special output
|
||||
// from the cache to meet specific requirements. It is frequently implemented in
|
||||
// conjunction with ScaledDotProductAttention.
|
||||
type BackendCacheConfig interface {
|
||||
CacheConfig() CacheConfig
|
||||
}
|
||||
|
||||
// CacheConfig controls optimizations (mostly backend-specific) that may transform
|
||||
// the output the cache to work better with specific kernels.
|
||||
type CacheConfig struct {
|
||||
// CachePadding specifies the multiple for the number of tokens of cache history
|
||||
// that will be returned from cache Get for k, v and mask. The capacity of the
|
||||
// cache itself will also be increased to a multiple of this size if needed.
|
||||
CachePadding int
|
||||
|
||||
// PermutedV performs Permute(ctx, 1, 2, 0, 3) on v tensors stored via Put
|
||||
// and return the permuted version via Get. This uses the cache copy operation
|
||||
// to avoid a Contiguous call on the permuted tensor.
|
||||
PermutedV bool
|
||||
|
||||
// MaskDType specifies the data type for generating the mask. If unset it will
|
||||
// default to DTypeF32.
|
||||
MaskDType DType
|
||||
}
|
||||
|
||||
// BackendParams controls how the backend loads and executes models
|
||||
type BackendParams struct {
|
||||
// AllocMemory causes the backend to allocate memory for the model. If
|
||||
// false, this is only being used for discovering the required amount of
|
||||
// memory and cannot load the model for running.
|
||||
AllocMemory bool
|
||||
|
||||
// NumThreads sets the number of threads to use if running on the CPU
|
||||
NumThreads int
|
||||
|
||||
// FlashAttention indicates that we should use a fused flash attention kernel
|
||||
FlashAttention FlashAttentionType
|
||||
}
|
||||
|
||||
var backends = make(map[string]func(string, BackendParams) (Backend, error))
|
||||
|
||||
func RegisterBackend(name string, f func(string, BackendParams) (Backend, error)) {
|
||||
if _, ok := backends[name]; ok {
|
||||
panic("backend: backend already registered")
|
||||
}
|
||||
|
||||
backends[name] = f
|
||||
}
|
||||
|
||||
func NewBackend(modelPath string, params BackendParams) (Backend, error) {
|
||||
if backend, ok := backends["ggml"]; ok {
|
||||
return backend(modelPath, params)
|
||||
}
|
||||
|
||||
return nil, fmt.Errorf("unsupported backend")
|
||||
}
|
||||
|
||||
type Context interface {
|
||||
Empty(dtype DType, shape ...int) Tensor
|
||||
Zeros(dtype DType, shape ...int) Tensor
|
||||
FromBytes(dtype DType, s []byte, shape ...int) Tensor
|
||||
FromFloats(s []float32, shape ...int) Tensor
|
||||
FromInts(s []int32, shape ...int) Tensor
|
||||
|
||||
// Arange creates a 1D tensor with values within an interval (start, stop] increased by step.
|
||||
Arange(start, stop, step float32, dtype DType) Tensor
|
||||
|
||||
Forward(...Tensor) Context
|
||||
|
||||
// SetBatchSize provides a hint on the batch size to optimize processing
|
||||
// Uses heuristics if not set
|
||||
SetBatchSize(int)
|
||||
|
||||
Compute(...Tensor)
|
||||
ComputeWithNotify(func(), ...Tensor) // notify callback once compute has begun
|
||||
|
||||
// Reserve is analogous to Compute but rather than executing a
|
||||
// graph, simply preallocates memory. Typically called with a
|
||||
// worst case graph to ensure all resources are available for
|
||||
// for future inference.
|
||||
Reserve()
|
||||
|
||||
MaxGraphNodes() int
|
||||
Close()
|
||||
|
||||
// Input returns a context appropriate for creating tensors that are
|
||||
// inputs to the model (which includes things like output locations)
|
||||
Input() Context
|
||||
|
||||
// Layer returns a context appropriate for creating intermediate tensors
|
||||
Layer(int) Context
|
||||
}
|
||||
|
||||
type Tensor interface {
|
||||
Dim(n int) int
|
||||
Stride(n int) int
|
||||
|
||||
Shape() []int
|
||||
DType() DType
|
||||
Cast(ctx Context, dtype DType) Tensor
|
||||
|
||||
Bytes() []byte
|
||||
Floats() []float32
|
||||
BackendGet() []float32
|
||||
|
||||
FromBytes([]byte)
|
||||
FromFloats([]float32)
|
||||
FromInts([]int32)
|
||||
|
||||
Add(ctx Context, t2 Tensor) Tensor
|
||||
Sub(ctx Context, t2 Tensor) Tensor
|
||||
Mul(ctx Context, t2 Tensor) Tensor
|
||||
Div(ctx Context, t2 Tensor) Tensor
|
||||
|
||||
Mulmat(ctx Context, t2 Tensor) Tensor
|
||||
MulmatFullPrec(ctx Context, t2 Tensor) Tensor
|
||||
MulmatID(ctx Context, t2, ids Tensor) Tensor
|
||||
AddID(ctx Context, t2, ids Tensor) Tensor
|
||||
|
||||
Softmax(ctx Context) Tensor
|
||||
L2Norm(ctx Context, eps float32) Tensor
|
||||
LayerNorm(ctx Context, weight, bias Tensor, eps float32) Tensor
|
||||
RMSNorm(ctx Context, weight Tensor, eps float32) Tensor
|
||||
Scale(ctx Context, s float64) Tensor
|
||||
SumRows(ctx Context) Tensor
|
||||
|
||||
AvgPool2D(ctx Context, k, s int, p float32) Tensor
|
||||
Conv2D(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor
|
||||
Conv3D(ctx Context, weight Tensor, c, s0, s1, s2, p0, p1, p2, d0, d1, d2 int) Tensor
|
||||
Conv1DDW(ctx Context, weight Tensor, s, p, d int) Tensor
|
||||
SSMConv(ctx Context, kernel Tensor) Tensor
|
||||
SSMScan(ctx Context, x, dt, A, B, C, ids Tensor) Tensor
|
||||
|
||||
IM2Col(ctx Context, weight Tensor, s0, s1, p0, p1, d0, d1 int) Tensor
|
||||
|
||||
Sin(ctx Context) Tensor
|
||||
Cos(ctx Context) Tensor
|
||||
Tanh(ctx Context) Tensor
|
||||
GELU(ctx Context, up ...Tensor) Tensor
|
||||
GELU_ERF(ctx Context) Tensor
|
||||
QuickGELU(ctx Context, up ...Tensor) Tensor
|
||||
SILU(ctx Context, up ...Tensor) Tensor
|
||||
RELU(ctx Context, up ...Tensor) Tensor
|
||||
Sigmoid(ctx Context) Tensor
|
||||
SigmoidOut(ctx Context) Tensor
|
||||
|
||||
// AlphaLimitSILU is a variant of SILU that clamps the input to the range [-limit, limit]
|
||||
SILUAlphaLimit(ctx Context, up Tensor, alpha, limit float32) Tensor
|
||||
|
||||
Reshape(ctx Context, shape ...int) Tensor
|
||||
View(ctx Context, offset int, shape ...int) Tensor
|
||||
Permute(ctx Context, shape ...int) Tensor
|
||||
Contiguous(ctx Context, shape ...int) Tensor
|
||||
|
||||
Pad(ctx Context, shape ...int) Tensor
|
||||
// PadExt pads with independent left/right amounts per dimension.
|
||||
// Arguments: lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3 for dims 0-3.
|
||||
PadExt(ctx Context, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3 int) Tensor
|
||||
|
||||
Stack(ctx Context, dim int, s ...Tensor) Tensor
|
||||
|
||||
// Repeat repeats the tensor n times along dimension dim
|
||||
Repeat(ctx Context, dim, n int) Tensor
|
||||
Concat(ctx Context, t2 Tensor, dim int) Tensor
|
||||
Rows(ctx Context, t2 Tensor) Tensor
|
||||
SetRows(ctx Context, src Tensor, idxs Tensor) Tensor
|
||||
SetInplace(ctx Context, src Tensor, nb1, nb2, nb3, offset int) Tensor
|
||||
Copy(ctx Context, t2 Tensor) Tensor
|
||||
Duplicate(ctx Context) Tensor
|
||||
|
||||
Slice(ctx Context, dim, low, high, step int) Tensor
|
||||
Chunk(ctx Context, dim int, size int) []Tensor
|
||||
ChunkSections(ctx Context, dim int, sections ...int) []Tensor
|
||||
|
||||
TopK(ctx Context, k int) Tensor
|
||||
Argsort(ctx Context) Tensor
|
||||
Mean(ctx Context) Tensor
|
||||
Variance(ctx Context) Tensor
|
||||
Stddev(ctx Context) Tensor
|
||||
Sqr(ctx Context) Tensor
|
||||
Sqrt(ctx Context) Tensor
|
||||
Exp(ctx Context) Tensor
|
||||
Neg(ctx Context) Tensor
|
||||
|
||||
// Clamp clamps values to [min, max] range
|
||||
Clamp(ctx Context, min, max float32) Tensor
|
||||
|
||||
// Softplus computes ln(1 + exp(x))
|
||||
Softplus(ctx Context) Tensor
|
||||
|
||||
// CumSum computes cumulative sum along dimension 0
|
||||
CumSum(ctx Context) Tensor
|
||||
|
||||
// Diag creates a diagonal matrix from a 1D tensor
|
||||
Diag(ctx Context) Tensor
|
||||
|
||||
// Tri converts a matrix to triangular form (0=upper+diag, 1=upper, 2=lower+diag, 3=lower)
|
||||
Tri(ctx Context, triType int) Tensor
|
||||
|
||||
// Fill fills a tensor with a constant value (in-place)
|
||||
Fill(ctx Context, value float32) Tensor
|
||||
|
||||
// Repeat4D repeats tensor to match target shape
|
||||
Repeat4D(ctx Context, dim0, dim1, dim2, dim3 int) Tensor
|
||||
|
||||
// SolveTri solves a triangular system Ax = B
|
||||
SolveTri(ctx Context, b Tensor, lower, left, unitDiag bool) Tensor
|
||||
|
||||
Interpolate(ctx Context, dims [4]int, samplingMode SamplingMode) Tensor
|
||||
}
|
||||
|
||||
// ScaledDotProductAttention implements a fused attention
|
||||
// operation equivalent to following code on a tensor named
|
||||
// query:
|
||||
//
|
||||
// query = query.Permute(ctx, 0, 2, 1, 3)
|
||||
// key = key.Permute(ctx, 0, 2, 1, 3)
|
||||
// value = value.Permute(ctx, 1, 2, 0, 3).Contiguous(ctx)
|
||||
//
|
||||
// kq := key.MulmatFullPrec(ctx, query)
|
||||
//
|
||||
// kq = kq.Scale(ctx, scale)
|
||||
//
|
||||
// if mask != nil {
|
||||
// kq = kq.Add(ctx, mask)
|
||||
// }
|
||||
//
|
||||
// kq = kq.Softmax(ctx)
|
||||
//
|
||||
// kqv := value.Mulmat(ctx, kq)
|
||||
// return kqv.Permute(ctx, 0, 2, 1, 3).Contiguous(ctx)
|
||||
//
|
||||
// cacheConfigApplied indicates whether the optimizations requested through CacheConfig have been performed
|
||||
type ScaledDotProductAttention interface {
|
||||
ScaledDotProductAttention(ctx Context, key, value, mask, sinks Tensor, vmla Tensor, scale float64, cacheConfigApplied bool) Tensor
|
||||
}
|
||||
|
||||
type number interface {
|
||||
~int | ~int8 | ~int16 | ~int32 | ~int64 |
|
||||
~uint | ~uint8 | ~uint16 | ~uint32 | ~uint64 |
|
||||
~float32 | ~float64 |
|
||||
~complex64 | ~complex128
|
||||
}
|
||||
|
||||
func mul[T number](s ...T) T {
|
||||
p := T(1)
|
||||
for _, v := range s {
|
||||
p *= v
|
||||
}
|
||||
|
||||
return p
|
||||
}
|
||||
|
||||
type DumpOptions func(*dumpOptions)
|
||||
|
||||
// DumpWithPrecision sets the number of decimal places to print. Applies to float32 and float64.
|
||||
func DumpWithPrecision(n int) DumpOptions {
|
||||
return func(opts *dumpOptions) {
|
||||
opts.Precision = n
|
||||
}
|
||||
}
|
||||
|
||||
// DumpWithThreshold sets the threshold for printing the entire tensor. If the number of elements
|
||||
// is less than or equal to this value, the entire tensor will be printed. Otherwise, only the
|
||||
// beginning and end of each dimension will be printed.
|
||||
func DumpWithThreshold(n int) DumpOptions {
|
||||
return func(opts *dumpOptions) {
|
||||
opts.Threshold = n
|
||||
}
|
||||
}
|
||||
|
||||
// DumpWithEdgeItems sets the number of elements to print at the beginning and end of each dimension.
|
||||
func DumpWithEdgeItems(n int) DumpOptions {
|
||||
return func(opts *dumpOptions) {
|
||||
opts.EdgeItems = n
|
||||
}
|
||||
}
|
||||
|
||||
type dumpOptions struct {
|
||||
Precision, Threshold, EdgeItems int
|
||||
}
|
||||
|
||||
func Dump(ctx Context, t Tensor, optsFuncs ...DumpOptions) string {
|
||||
opts := dumpOptions{Precision: 4, Threshold: 1000, EdgeItems: 3}
|
||||
for _, optsFunc := range optsFuncs {
|
||||
optsFunc(&opts)
|
||||
}
|
||||
|
||||
if mul(t.Shape()...) <= opts.Threshold {
|
||||
opts.EdgeItems = math.MaxInt
|
||||
}
|
||||
|
||||
switch t.DType() {
|
||||
case DTypeF32:
|
||||
return dump[[]float32](ctx, t, opts.EdgeItems, func(f float32) string {
|
||||
return strconv.FormatFloat(float64(f), 'f', opts.Precision, 32)
|
||||
})
|
||||
case DTypeF16, DTypeQ80, DTypeQ40:
|
||||
f32 := ctx.Input().Empty(DTypeF32, t.Shape()...)
|
||||
f32 = t.Copy(ctx, f32)
|
||||
return dump[[]float32](ctx, f32, opts.EdgeItems, func(f float32) string {
|
||||
return strconv.FormatFloat(float64(f), 'f', opts.Precision, 32)
|
||||
})
|
||||
case DTypeI32:
|
||||
return dump[[]int32](ctx, t, opts.EdgeItems, func(i int32) string {
|
||||
return strconv.FormatInt(int64(i), 10)
|
||||
})
|
||||
default:
|
||||
return "<unsupported>"
|
||||
}
|
||||
}
|
||||
|
||||
func dump[S ~[]E, E number](ctx Context, t Tensor, items int, fn func(E) string) string {
|
||||
if t.Bytes() == nil {
|
||||
ctx.Forward(t).Compute(t)
|
||||
}
|
||||
|
||||
s := make(S, mul(t.Shape()...))
|
||||
if err := binary.Read(bytes.NewBuffer(t.Bytes()), binary.LittleEndian, &s); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
|
||||
shape := t.Shape()
|
||||
slices.Reverse(shape)
|
||||
|
||||
var sb strings.Builder
|
||||
var f func([]int, int)
|
||||
f = func(dims []int, stride int) {
|
||||
prefix := strings.Repeat(" ", len(shape)-len(dims)+1)
|
||||
sb.WriteString("[")
|
||||
defer func() { sb.WriteString("]") }()
|
||||
for i := 0; i < dims[0]; i++ {
|
||||
if i >= items && i < dims[0]-items {
|
||||
sb.WriteString("..., ")
|
||||
// skip to next printable element
|
||||
skip := dims[0] - 2*items
|
||||
if len(dims) > 1 {
|
||||
stride += mul(append(dims[1:], skip)...)
|
||||
fmt.Fprint(&sb, strings.Repeat("\n", len(dims)-1), prefix)
|
||||
}
|
||||
i += skip - 1
|
||||
} else if len(dims) > 1 {
|
||||
f(dims[1:], stride)
|
||||
stride += mul(dims[1:]...)
|
||||
if i < dims[0]-1 {
|
||||
fmt.Fprint(&sb, ",", strings.Repeat("\n", len(dims)-1), prefix)
|
||||
}
|
||||
} else {
|
||||
text := fn(s[stride+i])
|
||||
if len(text) > 0 && text[0] != '-' {
|
||||
sb.WriteString(" ")
|
||||
}
|
||||
|
||||
sb.WriteString(text)
|
||||
if i < dims[0]-1 {
|
||||
sb.WriteString(", ")
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
f(shape, 0)
|
||||
|
||||
return sb.String()
|
||||
}
|
||||
|
||||
type DType int
|
||||
|
||||
const (
|
||||
DTypeOther DType = iota
|
||||
DTypeF32
|
||||
DTypeF16
|
||||
DTypeQ80
|
||||
DTypeQ40
|
||||
DTypeI32
|
||||
DTypeMXFP4
|
||||
)
|
||||
|
||||
type SamplingMode int
|
||||
|
||||
const (
|
||||
SamplingModeNearest SamplingMode = iota
|
||||
SamplingModeBilinear
|
||||
)
|
||||
-5
@@ -1,5 +0,0 @@
|
||||
// This file has been autogenerated by generate_cu_files.py, do not edit manually.
|
||||
|
||||
#include "../fattn-tile.cuh"
|
||||
|
||||
DECL_FATTN_TILE_CASE(512, 512);
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"runtime"
|
||||
"sort"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
@@ -271,31 +270,6 @@ func SimilarDeviceMemory(a, b uint64) bool {
|
||||
return maxMemory-min(a, b) <= tolerance
|
||||
}
|
||||
|
||||
// For a SameBackendDevice, return true if b is better than a
|
||||
// e.g. newer GPU library version
|
||||
func (a DeviceInfo) IsBetter(b DeviceInfo) bool {
|
||||
aLib := a.LibraryPath[len(a.LibraryPath)-1]
|
||||
bLib := b.LibraryPath[len(b.LibraryPath)-1]
|
||||
if aLib == bLib {
|
||||
return false
|
||||
}
|
||||
aLibSplit := strings.SplitN(aLib, "_", 2)
|
||||
bLibSplit := strings.SplitN(bLib, "_", 2)
|
||||
if len(aLibSplit) < 2 || len(bLibSplit) < 2 {
|
||||
return false
|
||||
}
|
||||
if aLibSplit[0] != bLibSplit[0] {
|
||||
slog.Debug("unexpected libraries", "a", aLib, "b", bLib)
|
||||
return false
|
||||
}
|
||||
if aLibSplit[1] == bLibSplit[1] {
|
||||
return false
|
||||
}
|
||||
cmp := []string{aLibSplit[1], bLibSplit[1]}
|
||||
sort.Sort(sort.Reverse(sort.StringSlice(cmp)))
|
||||
return cmp[0] == bLibSplit[1]
|
||||
}
|
||||
|
||||
// FlashAttentionSupported reports whether flash attention can be used across
|
||||
// all selected devices.
|
||||
func FlashAttentionSupported(l []DeviceInfo) bool {
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
package progress
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// StepBar displays step-based progress.
|
||||
type StepBar struct {
|
||||
message string
|
||||
current int
|
||||
total int
|
||||
}
|
||||
|
||||
func NewStepBar(message string, total int) *StepBar {
|
||||
return &StepBar{message: message, total: total}
|
||||
}
|
||||
|
||||
func (s *StepBar) Set(current int) {
|
||||
s.current = current
|
||||
}
|
||||
|
||||
func (s *StepBar) String() string {
|
||||
percent := float64(s.current) / float64(s.total) * 100
|
||||
barWidth := s.total
|
||||
empty := barWidth - s.current
|
||||
|
||||
// "Generating 0% ▕ ▏ 0/9"
|
||||
return fmt.Sprintf("%s %3.0f%% ▕%s%s▏ %d/%d",
|
||||
s.message, percent,
|
||||
strings.Repeat("█", s.current), strings.Repeat(" ", empty),
|
||||
s.current, s.total)
|
||||
}
|
||||
@@ -1,359 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"cmp"
|
||||
"fmt"
|
||||
"iter"
|
||||
"log/slog"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/dlclark/regexp2"
|
||||
heap "github.com/emirpasic/gods/v2/trees/binaryheap"
|
||||
"github.com/ollama/ollama/logutil"
|
||||
)
|
||||
|
||||
type BytePairEncoding struct {
|
||||
vocab *Vocabulary
|
||||
regexps []*regexp2.Regexp
|
||||
spaceToSpmSep bool // When true, normalize spaces to ▁ instead of GPT-2 byte-level encoding
|
||||
}
|
||||
|
||||
var _ Tokenizer = (*BytePairEncoding)(nil)
|
||||
|
||||
// BPEOption configures BytePairEncoding behavior
|
||||
type BPEOption func(*BytePairEncoding)
|
||||
|
||||
// WithSentencePieceNormalizer enables ▁ space normalization instead of GPT-2 byte-level encoding.
|
||||
func WithSentencePieceNormalizer() BPEOption {
|
||||
return func(bpe *BytePairEncoding) {
|
||||
bpe.spaceToSpmSep = true
|
||||
}
|
||||
}
|
||||
|
||||
func NewBytePairEncoding(vocab *Vocabulary, pretokenizer ...string) BytePairEncoding {
|
||||
return newBytePairEncoding(vocab, pretokenizer)
|
||||
}
|
||||
|
||||
func NewBytePairEncodingWithOptions(vocab *Vocabulary, pretokenizer []string, opts ...BPEOption) BytePairEncoding {
|
||||
bpe := newBytePairEncoding(vocab, pretokenizer, opts...)
|
||||
return bpe
|
||||
}
|
||||
|
||||
func newBytePairEncoding(vocab *Vocabulary, pretokenizer []string, opts ...BPEOption) BytePairEncoding {
|
||||
bpe := BytePairEncoding{
|
||||
vocab: vocab,
|
||||
}
|
||||
|
||||
for _, opt := range opts {
|
||||
opt(&bpe)
|
||||
}
|
||||
|
||||
if len(pretokenizer) == 0 && !bpe.spaceToSpmSep {
|
||||
// set default byte-level pretokenizer if none provided, e.g.
|
||||
// https://github.com/huggingface/tokenizer/blob/main/tokenizer/src/pre_tokenizer/byte_level.rs#L44
|
||||
pretokenizer = []string{`'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`}
|
||||
}
|
||||
|
||||
bpe.regexps = slices.Collect(func(yield func(*regexp2.Regexp) bool) {
|
||||
for _, p := range pretokenizer {
|
||||
if !yield(regexp2.MustCompile(p, regexp2.RE2)) {
|
||||
return
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
return bpe
|
||||
}
|
||||
|
||||
func (bpe *BytePairEncoding) split(s string) iter.Seq[string] {
|
||||
parts := []string{s}
|
||||
for _, re := range bpe.regexps {
|
||||
parts = slices.Collect(func(yield func(string) bool) {
|
||||
for _, part := range parts {
|
||||
r := []rune(part)
|
||||
var offset int
|
||||
for m, _ := re.FindRunesMatch(r); m != nil; m, _ = re.FindNextMatch(m) {
|
||||
if offset-m.Index != 0 {
|
||||
if !yield(string(r[offset:m.Index])) {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if !yield(m.String()) {
|
||||
return
|
||||
}
|
||||
|
||||
offset = m.Index + m.Length
|
||||
}
|
||||
|
||||
if offset < len(r) {
|
||||
if !yield(string(r[offset:])) {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
return slices.Values(parts)
|
||||
}
|
||||
|
||||
// fragment is a string fragment and their corresponding token IDs
|
||||
type fragment struct {
|
||||
value string
|
||||
ids []int32
|
||||
}
|
||||
|
||||
// pair is a pair of runes and its rank
|
||||
type pair struct {
|
||||
a, b int
|
||||
rank int
|
||||
value string
|
||||
}
|
||||
|
||||
type merge struct {
|
||||
p, n int
|
||||
runes []rune
|
||||
}
|
||||
|
||||
func (bpe BytePairEncoding) Encode(s string, addSpecial bool) ([]int32, error) {
|
||||
fragments := []fragment{{value: s}}
|
||||
for _, special := range bpe.vocab.SpecialVocabulary() {
|
||||
// TODO: process special tokens concurrently
|
||||
id := bpe.vocab.Encode(special)
|
||||
for i := 0; i < len(fragments); i++ {
|
||||
frag := fragments[i]
|
||||
if len(frag.ids) > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var middle []fragment
|
||||
switch i := strings.Index(frag.value, special); {
|
||||
case i < 0:
|
||||
middle = append(middle, frag)
|
||||
case i > 0:
|
||||
middle = append(middle, fragment{value: frag.value[:i]})
|
||||
fallthrough
|
||||
default:
|
||||
middle = append(middle, fragment{value: special, ids: []int32{id}})
|
||||
if rest := frag.value[i+len(special):]; rest != "" {
|
||||
middle = append(middle, fragment{value: rest})
|
||||
}
|
||||
}
|
||||
|
||||
fragments = append(fragments[:i], append(middle, fragments[i+1:]...)...)
|
||||
}
|
||||
}
|
||||
|
||||
var ids []int32
|
||||
for _, frag := range fragments {
|
||||
if len(frag.ids) > 0 {
|
||||
ids = append(ids, frag.ids...)
|
||||
continue
|
||||
}
|
||||
|
||||
for split := range bpe.split(frag.value) {
|
||||
// TODO: process splits concurrently
|
||||
var normalized string
|
||||
if bpe.spaceToSpmSep {
|
||||
// SentencePiece-style: replace spaces with ▁
|
||||
normalized = strings.ReplaceAll(split, " ", spmWhitespaceSep)
|
||||
} else {
|
||||
// GPT-2 byte-level: map bytes to shifted Unicode codepoints
|
||||
var sb strings.Builder
|
||||
for _, b := range []byte(split) {
|
||||
r := rune(b)
|
||||
switch {
|
||||
case r == 0x00ad:
|
||||
r = 0x0143
|
||||
case r <= 0x0020:
|
||||
r = r + 0x0100
|
||||
case r >= 0x007f && r <= 0x00a0:
|
||||
r = r + 0x00a2
|
||||
}
|
||||
sb.WriteRune(r)
|
||||
}
|
||||
normalized = sb.String()
|
||||
}
|
||||
|
||||
// short circuit if the fragment is in the vocabulary
|
||||
if id := bpe.vocab.Encode(normalized); id >= 0 {
|
||||
ids = append(ids, id)
|
||||
continue
|
||||
}
|
||||
|
||||
runes := []rune(normalized)
|
||||
merges := make([]merge, len(runes))
|
||||
for r := range runes {
|
||||
merges[r] = merge{
|
||||
p: r - 1,
|
||||
n: r + 1,
|
||||
runes: []rune{runes[r]},
|
||||
}
|
||||
}
|
||||
|
||||
pairwise := func(a, b int) *pair {
|
||||
if a < 0 || b >= len(runes) {
|
||||
return nil
|
||||
}
|
||||
|
||||
left, right := string(merges[a].runes), string(merges[b].runes)
|
||||
rank := bpe.vocab.Merge(left, right)
|
||||
if rank < 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
return &pair{
|
||||
a: a,
|
||||
b: b,
|
||||
rank: rank,
|
||||
value: left + right,
|
||||
}
|
||||
}
|
||||
|
||||
pairs := heap.NewWith(func(i, j *pair) int {
|
||||
return cmp.Compare(i.rank, j.rank)
|
||||
})
|
||||
|
||||
for i := range len(runes) - 1 {
|
||||
if pair := pairwise(i, i+1); pair != nil {
|
||||
pairs.Push(pair)
|
||||
}
|
||||
}
|
||||
|
||||
for !pairs.Empty() {
|
||||
pair, _ := pairs.Pop()
|
||||
|
||||
left, right := merges[pair.a], merges[pair.b]
|
||||
if len(left.runes) == 0 || len(right.runes) == 0 ||
|
||||
string(left.runes)+string(right.runes) != pair.value {
|
||||
continue
|
||||
}
|
||||
|
||||
if id := bpe.vocab.Encode(pair.value); id < 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
merges[pair.a].runes = append(left.runes, right.runes...)
|
||||
merges[pair.b].runes = nil
|
||||
|
||||
merges[pair.a].n = right.n
|
||||
if right.n < len(merges) {
|
||||
merges[right.n].p = pair.a
|
||||
}
|
||||
|
||||
if pair := pairwise(merges[pair.a].p, pair.a); pair != nil {
|
||||
pairs.Push(pair)
|
||||
}
|
||||
|
||||
if pair := pairwise(pair.a, merges[pair.a].n); pair != nil {
|
||||
pairs.Push(pair)
|
||||
}
|
||||
}
|
||||
|
||||
for _, merge := range merges {
|
||||
if len(merge.runes) > 0 {
|
||||
if id := bpe.vocab.Encode(string(merge.runes)); id >= 0 {
|
||||
ids = append(ids, id)
|
||||
} else if bpe.spaceToSpmSep {
|
||||
// SentencePiece byte fallback: encode each UTF-8 byte as <0xHH>
|
||||
for _, b := range []byte(string(merge.runes)) {
|
||||
if id := bpe.vocab.Encode(fmt.Sprintf("<0x%02X>", b)); id >= 0 {
|
||||
ids = append(ids, id)
|
||||
} else {
|
||||
slog.Debug("unknown byte token", "byte", b)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if addSpecial {
|
||||
ids = bpe.vocab.addSpecials(ids)
|
||||
}
|
||||
|
||||
logutil.Trace("encoded", "string", s, "ids", ids)
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
type lazyIdsString struct {
|
||||
ids []int32
|
||||
}
|
||||
|
||||
func (l lazyIdsString) LogValue() slog.Value {
|
||||
return slog.AnyValue(fmt.Sprint(l.ids))
|
||||
}
|
||||
|
||||
func (bpe BytePairEncoding) Decode(ids []int32) (string, error) {
|
||||
var sb strings.Builder
|
||||
|
||||
// SentencePiece-style BPE stores true Unicode codepoints in the vocab
|
||||
// (plus ▁ as a whitespace marker), so decoding should pass runes through
|
||||
// directly instead of applying the GPT-2 byte-level reverse mapping.
|
||||
// Without this, codepoints in the 0x0100-0x0142 range (e.g. ą ę ć ł)
|
||||
// get mangled by the GPT-2 reversal into control characters.
|
||||
if bpe.spaceToSpmSep {
|
||||
for _, id := range ids {
|
||||
data := bpe.vocab.Decode(id)
|
||||
|
||||
// SentencePiece byte tokens: "<0xHH>" → raw byte
|
||||
if len(data) == 6 && strings.HasPrefix(data, "<0x") && strings.HasSuffix(data, ">") {
|
||||
if b, err := strconv.ParseUint(data[3:5], 16, 8); err == nil {
|
||||
sb.WriteByte(byte(b))
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
for _, r := range data {
|
||||
if r == 0x2581 { // ▁ (LOWER ONE EIGHTH BLOCK)
|
||||
sb.WriteByte(' ')
|
||||
} else {
|
||||
sb.WriteRune(r)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logutil.Trace("decoded", "string", sb.String(), "from", lazyIdsString{ids: ids})
|
||||
return sb.String(), nil
|
||||
}
|
||||
|
||||
for _, id := range ids {
|
||||
for _, r := range bpe.vocab.Decode(id) {
|
||||
// GPT-2 byte-level BPE uses Unicode chars in the 0x0100-0x0143
|
||||
// range to represent bytes. Remap them back to actual bytes.
|
||||
switch {
|
||||
case r == 0x0100:
|
||||
// this produces 0x00 aka NULL
|
||||
continue
|
||||
case r == 0x0143:
|
||||
r = 0x00ad
|
||||
case r > 0x0100 && r <= 0x0120:
|
||||
r = r - 0x0100
|
||||
case r > 0x0120 && r <= 0x0142:
|
||||
r = r - 0x00a2
|
||||
case r > 0x0143:
|
||||
// Non-GPT2 rune (e.g., SentencePiece-style BPE).
|
||||
// Handle ▁ as word separator, otherwise write the rune as-is.
|
||||
if r == 0x2581 { // ▁ (LOWER ONE EIGHTH BLOCK)
|
||||
sb.WriteByte(' ')
|
||||
} else {
|
||||
sb.WriteRune(r)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
// NOTE: not using WriteRune here because it writes the UTF-8
|
||||
// encoding of the rune which is _not_ what we want
|
||||
if err := sb.WriteByte(byte(r)); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logutil.Trace("decoded", "string", sb.String(), "from", lazyIdsString{ids: ids})
|
||||
return sb.String(), nil
|
||||
}
|
||||
@@ -1,639 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"math"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
)
|
||||
|
||||
func llama(t testing.TB) BytePairEncoding {
|
||||
t.Helper()
|
||||
|
||||
f, err := os.Open(filepath.FromSlash("testdata/llama3.2/encoder.json"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
vocab := make(map[string]int32)
|
||||
if err := json.NewDecoder(f).Decode(&vocab); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
types := make([]int32, len(vocab))
|
||||
tokens := make([]string, len(vocab))
|
||||
for token, id := range vocab {
|
||||
tokens[id] = token
|
||||
types[id] = 1
|
||||
}
|
||||
|
||||
for _, token := range []string{"<|begin_of_text|>", "<|end_of_text|>"} {
|
||||
if _, ok := vocab[token]; !ok {
|
||||
tokens = append(tokens, token) //nolint:makezero
|
||||
types = append(types, 3) //nolint:makezero
|
||||
vocab[token] = int32(len(vocab))
|
||||
}
|
||||
}
|
||||
|
||||
f, err = os.Open(filepath.FromSlash("testdata/llama3.2/vocab.bpe"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer f.Close()
|
||||
|
||||
merges := make([]string, 0, 50000)
|
||||
|
||||
scanner := bufio.NewScanner(f)
|
||||
for scanner.Scan() {
|
||||
if !strings.HasPrefix(scanner.Text(), "#") {
|
||||
merges = append(merges, scanner.Text())
|
||||
}
|
||||
}
|
||||
|
||||
return NewBytePairEncoding(
|
||||
&Vocabulary{
|
||||
Values: tokens,
|
||||
Types: types,
|
||||
Merges: merges,
|
||||
},
|
||||
"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
|
||||
)
|
||||
}
|
||||
|
||||
func TestLlama(t *testing.T) {
|
||||
tokenizer := llama(t)
|
||||
|
||||
t.Run("simple", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ids, err := tokenizer.Encode("hello world", true)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff([]int32{15339, 1917}, ids); diff != "" {
|
||||
t.Errorf("no match (-theirs +ours):\n%s", diff)
|
||||
}
|
||||
|
||||
s, err := tokenizer.Decode([]int32{15339, 1917})
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if s != "hello world" {
|
||||
t.Errorf("got %q, want hello world", s)
|
||||
}
|
||||
|
||||
ids, err = tokenizer.Encode("hello <|end_of_text|>", true)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff([]int32{15339, 220, 128001}, ids); diff != "" {
|
||||
t.Errorf("no match (-theirs +ours):\n%s", diff)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("simple repeated", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string][]int32{
|
||||
strings.Repeat("0", 1): {15},
|
||||
strings.Repeat("0", 2): {410},
|
||||
strings.Repeat("0", 3): {931},
|
||||
strings.Repeat("0", 4): {931, 15},
|
||||
strings.Repeat("0", 5): {931, 410},
|
||||
strings.Repeat("0", 6): {931, 931},
|
||||
strings.Repeat("0", 7): {931, 931, 15},
|
||||
strings.Repeat("0", 8): {931, 931, 410},
|
||||
strings.Repeat("0", 9): {931, 931, 931},
|
||||
strings.Repeat("0", 10): {931, 931, 931, 15},
|
||||
strings.Repeat("0", 11): {931, 931, 931, 410},
|
||||
strings.Repeat("0", 12): {931, 931, 931, 931},
|
||||
strings.Repeat("0", 13): {931, 931, 931, 931, 15},
|
||||
strings.Repeat("0", 14): {931, 931, 931, 931, 410},
|
||||
strings.Repeat("0", 15): {931, 931, 931, 931, 931},
|
||||
strings.Repeat("0", 16): {931, 931, 931, 931, 931, 15},
|
||||
strings.Repeat("0", 17): {931, 931, 931, 931, 931, 410},
|
||||
}
|
||||
|
||||
for s, want := range cases {
|
||||
ids, err := tokenizer.Encode(s, true)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(want, ids); diff != "" {
|
||||
t.Errorf("%q no match (-theirs +ours):\n%s", s, diff)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("basic roundtrip", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []string{
|
||||
"hello",
|
||||
"hello ",
|
||||
"hello ",
|
||||
" hello",
|
||||
" hello ",
|
||||
" hello ",
|
||||
"hello world",
|
||||
"请考试我的软件！12345",
|
||||
}
|
||||
|
||||
for _, want := range cases {
|
||||
ids, err := tokenizer.Encode(want, true)
|
||||
if err != nil {
|
||||
t.Error(err)
|
||||
}
|
||||
|
||||
if got, err := tokenizer.Decode(ids); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if got != want {
|
||||
t.Errorf("got %q, want %q", got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("special", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string][]int32{
|
||||
"<|begin_of_text|>A B!": {128000, 32, 426, 0},
|
||||
"<|begin_of_text|>A<|end_of_text|>B!": {128000, 32, 128001, 33, 0},
|
||||
"<|begin_of_text|>A<|end_of_text|>B<|begin_of_text|>!": {128000, 32, 128001, 33, 128000, 0},
|
||||
"<|begin_of_text|>A<|end_of_text|>B<|begin_of_text|>!<|end_of_text|>": {128000, 32, 128001, 33, 128000, 0, 128001},
|
||||
}
|
||||
|
||||
for s, want := range cases {
|
||||
ids, err := tokenizer.Encode(s, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff(want, ids); diff != "" {
|
||||
t.Errorf("no match (-theirs +ours):\n%s", diff)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("split", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string][]string{
|
||||
"Hello World!": {"Hello", " World", "!"},
|
||||
"I'm don't won't": {"I", "'m", " don", "'t", " won", "'t"},
|
||||
"In 2024 there are 366 days": {"In", " ", "202", "4", " there", " are", " ", "366", " days"},
|
||||
"Hello!! ...world": {"Hello", "!!", " ...", "world"},
|
||||
"Hello World": {"Hello", " ", " World"},
|
||||
"Hello\nWorld": {"Hello", "\n", "World"},
|
||||
"Hello, WORLD!! How's it going?": {"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?"},
|
||||
}
|
||||
|
||||
for s, want := range cases {
|
||||
got := slices.Collect(tokenizer.split(s))
|
||||
if diff := cmp.Diff(want, got); diff != "" {
|
||||
t.Errorf("no match (-theirs +ours):\n%s", diff)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("roundtriping 0x00-0xFF", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for b := 0x00; b <= 0xFF; b++ {
|
||||
input := string(rune(b))
|
||||
ids, err := tokenizer.Encode(input, false)
|
||||
if err != nil {
|
||||
t.Errorf("failed to encode rune 0x%02X: %v", b, err)
|
||||
continue
|
||||
}
|
||||
|
||||
decoded, err := tokenizer.Decode(ids)
|
||||
if err != nil {
|
||||
t.Errorf("failed to decode rune 0x%02X: %v", b, err)
|
||||
continue
|
||||
}
|
||||
|
||||
if b == 0x00 {
|
||||
if len(decoded) != 0 {
|
||||
t.Errorf("Decode(Encode(0x00)) should be empty, got %v", ids)
|
||||
}
|
||||
continue
|
||||
}
|
||||
|
||||
if decoded != input {
|
||||
t.Errorf("rune 0x%02X failed roundtrip: got %q, want %q", b, decoded, input)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// spmBPE builds a SentencePiece-style BPE tokenizer for testing.
|
||||
//
|
||||
// Models that use SentencePiece BPE differ from GPT-2 BPE in how they
|
||||
// handle spaces: the vocabulary stores ▁ (U+2581) instead of GPT-2's
|
||||
// shifted-byte encoding (0x0100–0x0143). Without WithSentencePieceNormalizer,
|
||||
// spaces are mapped through the GPT-2 byte table which produces wrong token
|
||||
// IDs for any vocabulary that uses ▁-prefixed tokens. The decode path has
|
||||
// the inverse problem: high codepoints like CJK characters and ▁ itself
|
||||
// would be mangled by the GPT-2 reverse mapping instead of being passed
|
||||
// through (or converted to spaces in the ▁ case).
|
||||
func spmBPE(t testing.TB) BytePairEncoding {
|
||||
t.Helper()
|
||||
|
||||
tokens := []string{
|
||||
// Control tokens (low IDs, as in real SentencePiece vocabs)
|
||||
"<pad>", // 0
|
||||
"<eos>", // 1
|
||||
"<bos>", // 2
|
||||
"<|start>", // 3 - asymmetric open/close special tokens
|
||||
"<end|>", // 4
|
||||
"<|q>", // 5 - short special token (like <|"|>)
|
||||
|
||||
// ▁-prefixed word tokens (the core of what SPM BPE changes)
|
||||
"▁hello", // 6
|
||||
"▁world", // 7
|
||||
"hello", // 8
|
||||
"▁Run", // 9
|
||||
"▁a", // 10
|
||||
|
||||
// Punctuation and structure
|
||||
",", // 11
|
||||
"!", // 12
|
||||
":", // 13
|
||||
"{", // 14
|
||||
"}", // 15
|
||||
|
||||
// Whitespace separator
|
||||
"▁", // 16
|
||||
|
||||
// Subword tokens used in tool-declaration-like patterns
|
||||
"description", // 17
|
||||
"▁command", // 18
|
||||
"declaration", // 19
|
||||
|
||||
// Unicode token for decode passthrough testing (must be > U+0143
|
||||
// to exercise the SPM decode path rather than GPT-2 byte reversal)
|
||||
"▁中文", // 20
|
||||
|
||||
// Unicode tokens with codepoints in the GPT-2 byte range (0x0100-0x0142).
|
||||
// Without the SPM decode path, these get mangled by GPT-2 byte reversal.
|
||||
"ą", // 21 (U+0105) — would become 0x05 via GPT-2 reversal
|
||||
"ę", // 22 (U+0119) — would become 0x19
|
||||
"ć", // 23 (U+0107) — would become 0x07
|
||||
"ł", // 24 (U+0142) — would become 0xA0
|
||||
|
||||
// Byte fallback tokens (SentencePiece BYTE type)
|
||||
"<0x00>", // 25
|
||||
"<0x01>", // 26
|
||||
}
|
||||
|
||||
// Add all 256 byte tokens starting at index 27
|
||||
for b := 2; b < 256; b++ {
|
||||
tokens = append(tokens, fmt.Sprintf("<0x%02X>", b))
|
||||
}
|
||||
|
||||
types := make([]int32, len(tokens))
|
||||
for i := range types {
|
||||
types[i] = TOKEN_TYPE_NORMAL
|
||||
}
|
||||
types[0] = TOKEN_TYPE_CONTROL // <pad>
|
||||
types[1] = TOKEN_TYPE_CONTROL // <eos>
|
||||
types[2] = TOKEN_TYPE_CONTROL // <bos>
|
||||
types[3] = TOKEN_TYPE_USER_DEFINED // <|start>
|
||||
types[4] = TOKEN_TYPE_USER_DEFINED // <end|>
|
||||
types[5] = TOKEN_TYPE_USER_DEFINED // <|q>
|
||||
for i := 21; i < len(types); i++ {
|
||||
types[i] = TOKEN_TYPE_BYTE
|
||||
}
|
||||
|
||||
return NewBytePairEncodingWithOptions(
|
||||
&Vocabulary{
|
||||
Values: tokens,
|
||||
Types: types,
|
||||
BOS: []int32{2},
|
||||
EOS: []int32{1},
|
||||
AddBOS: false,
|
||||
},
|
||||
// Empty pretokenizer list: falls back to the default pattern.
|
||||
// Real SentencePiece BPE models are configured this way.
|
||||
[]string{},
|
||||
WithSentencePieceNormalizer(),
|
||||
)
|
||||
}
|
||||
|
||||
func TestSentencePieceBPE(t *testing.T) {
|
||||
tok := spmBPE(t)
|
||||
|
||||
// Test 1: Space-to-▁ normalization and roundtrip.
|
||||
//
|
||||
// SentencePiece BPE has no pretokenizer — the BPE merges handle word
|
||||
// boundaries via ▁ markers. With no merges in the test vocab, multi-char
|
||||
// tokens won't be found, but the roundtrip must still be lossless.
|
||||
t.Run("spm space normalization roundtrip", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
for _, input := range []string{
|
||||
"hello",
|
||||
" hello",
|
||||
"hello, world!",
|
||||
" leading spaces",
|
||||
"multiple spaces",
|
||||
} {
|
||||
ids, err := tok.Encode(input, false)
|
||||
if err != nil {
|
||||
t.Fatalf("Encode(%q): %v", input, err)
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
t.Fatalf("Encode(%q) returned empty IDs", input)
|
||||
}
|
||||
|
||||
got, err := tok.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode(%v): %v", ids, err)
|
||||
}
|
||||
if got != input {
|
||||
t.Errorf("roundtrip %q: Decode(Encode) = %q", input, got)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
// Test 2: Special tokens interleaved with SPM-normalized text.
|
||||
//
|
||||
// This mimics tool declaration patterns like:
|
||||
// <|tool>declaration:bash{description:<|"|>Run a command<|"|>}<tool|>
|
||||
// where special tokens (<|tool>, <|"|>, <tool|>) must be extracted
|
||||
// first, then the remaining text fragments go through SPM normalization.
|
||||
t.Run("special tokens with spm text fragments", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
input := "<|start>declaration:description:<|q> Run a command<|q>}<end|>"
|
||||
ids, err := tok.Encode(input, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
// Special tokens should be extracted as single IDs at the right positions.
|
||||
// The text between them is SPM-normalized and BPE-encoded (specific IDs
|
||||
// depend on merges, so we verify the special token positions + roundtrip).
|
||||
specialPositions := map[int32]bool{3: true, 4: true, 5: true} // <|start>, <end|>, <|q>
|
||||
foundSpecials := 0
|
||||
for _, id := range ids {
|
||||
if specialPositions[id] {
|
||||
foundSpecials++
|
||||
}
|
||||
}
|
||||
if foundSpecials != 4 { // <|start>, <|q>, <|q>, <end|>
|
||||
t.Errorf("expected 4 special tokens, found %d in %v", foundSpecials, ids)
|
||||
}
|
||||
|
||||
// First token must be <|start>(3), last must be <end|>(4)
|
||||
if ids[0] != 3 {
|
||||
t.Errorf("first token = %d, want 3 (<|start>)", ids[0])
|
||||
}
|
||||
if ids[len(ids)-1] != 4 {
|
||||
t.Errorf("last token = %d, want 4 (<end|>)", ids[len(ids)-1])
|
||||
}
|
||||
})
|
||||
|
||||
// Test 3: Byte fallback for characters not in the vocabulary.
|
||||
//
|
||||
// SentencePiece vocabs include <0xHH> byte tokens for every byte value.
|
||||
// When a character (e.g. "ą" = U+0105 = C4 85) isn't in the vocab as a
|
||||
// direct token, the encoder must fall back to its UTF-8 bytes:
|
||||
// <0xC4> <0x85>. Without this fallback, the character is silently dropped.
|
||||
// See: https://github.com/ollama/ollama/issues/15229
|
||||
t.Run("byte fallback for unknown chars", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// "ą" is not in the vocab — should fall back to byte tokens
|
||||
ids, err := tok.Encode("ą", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Encode(ą): %v", err)
|
||||
}
|
||||
if len(ids) == 0 {
|
||||
t.Fatal("Encode(ą) returned empty IDs — character was silently dropped")
|
||||
}
|
||||
|
||||
got, err := tok.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode: %v", err)
|
||||
}
|
||||
if got != "ą" {
|
||||
t.Errorf("roundtrip = %q, want %q", got, "ą")
|
||||
}
|
||||
})
|
||||
|
||||
// Test 4: Byte fallback preserves known tokens around unknown chars.
|
||||
t.Run("byte fallback mixed with known tokens", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// "hello" is in vocab, "é" is not
|
||||
ids, err := tok.Encode("helloé", false)
|
||||
if err != nil {
|
||||
t.Fatalf("Encode: %v", err)
|
||||
}
|
||||
|
||||
got, err := tok.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode: %v", err)
|
||||
}
|
||||
if got != "helloé" {
|
||||
t.Errorf("roundtrip = %q, want %q", got, "helloé")
|
||||
}
|
||||
})
|
||||
|
||||
// Test 5: Decode doesn't mangle Unicode in the GPT-2 byte range.
|
||||
//
|
||||
// Characters like ą (U+0105), ę (U+0119), ć (U+0107), ł (U+0142) have
|
||||
// codepoints in the 0x0100-0x0142 range that GPT-2 byte reversal would
|
||||
// remap to control characters. SentencePiece decode must pass them through.
|
||||
t.Run("decode unicode in gpt2 byte range", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
// Token IDs 21-24 are ą, ę, ć, ł
|
||||
ids := []int32{21, 22, 23, 24}
|
||||
got, err := tok.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode: %v", err)
|
||||
}
|
||||
if got != "ąęćł" {
|
||||
t.Errorf("Decode = %q, want %q", got, "ąęćł")
|
||||
}
|
||||
})
|
||||
|
||||
// Test 6: Decode handles non-GPT2 Unicode correctly.
|
||||
//
|
||||
// GPT-2 BPE decode reverses the byte→codepoint shift for runes in
|
||||
// 0x0100–0x0143. But SentencePiece vocabs store real Unicode (CJK,
|
||||
// accented chars, etc.) which have codepoints well above 0x0143.
|
||||
// Without the > 0x0143 passthrough in Decode, these would be mangled
|
||||
// by the GPT-2 reverse mapping (e.g., written as raw bytes instead
|
||||
// of the original characters).
|
||||
t.Run("decode non-gpt2 unicode passthrough", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := map[string][]int32{
|
||||
" 中文": {20}, // ▁→space, then CJK passes through as-is
|
||||
}
|
||||
|
||||
for want, ids := range cases {
|
||||
got, err := tok.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatalf("Decode(%v): %v", ids, err)
|
||||
}
|
||||
if got != want {
|
||||
t.Errorf("Decode(%v) = %q, want %q", ids, got, want)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func BenchmarkBytePairEncoding(b *testing.B) {
|
||||
tokenizer := llama(b)
|
||||
bts, err := os.ReadFile(filepath.Join("testdata", "war-and-peace.txt"))
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
|
||||
for i := range 8 {
|
||||
n := min(int(math.Pow10(i)), len(bts))
|
||||
bts := bts[:n]
|
||||
b.Run("encode"+strconv.Itoa(n), func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
_, err := tokenizer.Encode(string(bts), true)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("decode"+strconv.Itoa(n), func(b *testing.B) {
|
||||
ids, err := tokenizer.Encode(string(bts), true)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
_, err := tokenizer.Decode(ids)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
b.Run("split"+strconv.Itoa(n), func(b *testing.B) {
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
slices.Collect(tokenizer.split(string(bts)))
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestBytePairEncodingSplitMultipleRegexpsPreservesOffsets(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
bpe := NewBytePairEncoding(
|
||||
nil,
|
||||
`(?:\r?\n)+(?!\r?\n)`,
|
||||
`(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+`,
|
||||
)
|
||||
|
||||
input := "One line\nTwo lines\n\nThree"
|
||||
got := slices.Collect(bpe.split(input))
|
||||
want := []string{"One", " line", "\n", "Two", " lines", "\n\n", "Three"}
|
||||
|
||||
if diff := cmp.Diff(want, got); diff != "" {
|
||||
t.Fatalf("split mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBytePairEncodingSplitRefactPreservesOffsets(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
bpe := NewBytePairEncoding(
|
||||
nil,
|
||||
`\p{N}`,
|
||||
`'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+`,
|
||||
)
|
||||
|
||||
input := "One line\nTwo lines\n\nThree"
|
||||
got := slices.Collect(bpe.split(input))
|
||||
want := []string{"One", " line", "\n", "Two", " lines", "\n", "\n", "Three"}
|
||||
|
||||
if diff := cmp.Diff(want, got); diff != "" {
|
||||
t.Fatalf("split mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBytePairEncodingSplitDeepSeekV3PreservesOffsets(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
bpe := NewBytePairEncoding(
|
||||
nil,
|
||||
"\\p{N}{1,3}",
|
||||
`[一-龥-ゟ゠-ヿ]+`,
|
||||
"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\\r\\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
|
||||
)
|
||||
|
||||
input := "One line\nTwo lines\n\nThree"
|
||||
got := slices.Collect(bpe.split(input))
|
||||
want := []string{"One", " line", "\n", "Two", " lines", "\n\n", "Three"}
|
||||
|
||||
if diff := cmp.Diff(want, got); diff != "" {
|
||||
t.Fatalf("split mismatch (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSplit(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
patterns,
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
name: "default",
|
||||
want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " 123", " 一二三"},
|
||||
},
|
||||
{
|
||||
name: "unicode",
|
||||
patterns: []string{
|
||||
"\\p{N}{1,3}",
|
||||
`[一-龥-ゟ゠-ヿ]+`,
|
||||
"[!\"#$%&'()*+,\\-./:;<=>?@\\[\\\\\\]^_`{|}~][A-Za-z]+|[^\r\n\\p{L}\\p{P}\\p{S}]?[\\p{L}\\p{M}]+| ?[\\p{P}\\p{S}]+[\r\n]*|\\s*[\r\n]+|\\s+(?!\\S)|\\s+",
|
||||
},
|
||||
want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " ", "123", " ", "一二三"},
|
||||
},
|
||||
{
|
||||
name: "individual digits",
|
||||
patterns: []string{
|
||||
"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
|
||||
},
|
||||
want: []string{"Hello", ",", " WORLD", "!!", " How", "'s", " it", " going", "?", " ", "1", "2", "3", " 一二三"},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range cases {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
tokenizer := NewBytePairEncoding(nil, tt.patterns...)
|
||||
if diff := cmp.Diff(tt.want, slices.Collect(tokenizer.split("Hello, WORLD!! How's it going? 123 一二三"))); diff != "" {
|
||||
t.Errorf("no match (-theirs +ours):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,241 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"container/heap"
|
||||
"fmt"
|
||||
"log/slog"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/ollama/ollama/logutil"
|
||||
)
|
||||
|
||||
const spmWhitespaceSep = "▁"
|
||||
|
||||
type SentencePiece struct {
|
||||
maxTokenLen int
|
||||
vocab *Vocabulary
|
||||
}
|
||||
|
||||
var _ Tokenizer = (*SentencePiece)(nil)
|
||||
|
||||
func NewSentencePiece(vocab *Vocabulary) SentencePiece {
|
||||
logutil.Trace("Tokens", "num tokens", len(vocab.Values), "vals", vocab.Values[:5], "scores", vocab.Scores[:5], "types", vocab.Types[:5])
|
||||
|
||||
counter := map[int]int{}
|
||||
var maxTokenLen int
|
||||
for cnt := range vocab.Types {
|
||||
switch vocab.Types[cnt] {
|
||||
case TOKEN_TYPE_NORMAL, TOKEN_TYPE_USER_DEFINED, TOKEN_TYPE_UNUSED:
|
||||
maxTokenLen = max(maxTokenLen, len(vocab.Values[cnt]))
|
||||
fallthrough
|
||||
default:
|
||||
counter[int(vocab.Types[cnt])] += 1
|
||||
}
|
||||
}
|
||||
|
||||
logutil.Trace("Token counts", "normal", counter[TOKEN_TYPE_NORMAL], "unknown", counter[TOKEN_TYPE_UNKNOWN], "control", counter[TOKEN_TYPE_CONTROL],
|
||||
"user defined", counter[TOKEN_TYPE_USER_DEFINED], "unused", counter[TOKEN_TYPE_UNUSED], "byte", counter[TOKEN_TYPE_BYTE],
|
||||
"max token len", maxTokenLen)
|
||||
|
||||
return SentencePiece{
|
||||
maxTokenLen: maxTokenLen,
|
||||
vocab: vocab,
|
||||
}
|
||||
}
|
||||
|
||||
func (spm SentencePiece) Encode(s string, addSpecial bool) ([]int32, error) {
|
||||
fragments := []fragment{{value: s}}
|
||||
for _, special := range spm.vocab.SpecialVocabulary() {
|
||||
id := spm.vocab.Encode(special)
|
||||
for i := 0; i < len(fragments); i++ {
|
||||
frag := fragments[i]
|
||||
if len(frag.ids) > 0 {
|
||||
continue
|
||||
}
|
||||
|
||||
var middle []fragment
|
||||
switch i := strings.Index(frag.value, special); {
|
||||
case i < 0:
|
||||
middle = append(middle, frag)
|
||||
case i > 0:
|
||||
middle = append(middle, fragment{value: frag.value[:i]})
|
||||
fallthrough
|
||||
default:
|
||||
middle = append(middle, fragment{value: special, ids: []int32{id}})
|
||||
if rest := frag.value[i+len(special):]; rest != "" {
|
||||
middle = append(middle, fragment{value: rest})
|
||||
}
|
||||
}
|
||||
|
||||
fragments = append(fragments[:i], append(middle, fragments[i+1:]...)...)
|
||||
}
|
||||
}
|
||||
|
||||
var ids []int32
|
||||
for _, frag := range fragments {
|
||||
if len(frag.ids) > 0 {
|
||||
ids = append(ids, frag.ids...)
|
||||
continue
|
||||
}
|
||||
|
||||
text := strings.ReplaceAll(frag.value, " ", spmWhitespaceSep)
|
||||
|
||||
if id := spm.vocab.Encode(text); id >= 0 {
|
||||
ids = append(ids, id)
|
||||
continue
|
||||
}
|
||||
|
||||
q := &queue{}
|
||||
heap.Init(q)
|
||||
|
||||
runes := []rune(text)
|
||||
merges := make([]merge, len(runes))
|
||||
for r := range runes {
|
||||
merges[r] = merge{
|
||||
p: r - 1,
|
||||
n: r + 1,
|
||||
runes: []rune{runes[r]},
|
||||
}
|
||||
}
|
||||
|
||||
pairwise := func(a, b int) *candidate {
|
||||
if a < 0 || b >= len(runes) {
|
||||
return nil
|
||||
}
|
||||
|
||||
left, right := string(merges[a].runes), string(merges[b].runes)
|
||||
if id := spm.vocab.Encode(left + right); id >= 0 {
|
||||
return &candidate{
|
||||
a: a,
|
||||
b: b,
|
||||
score: spm.vocab.Scores[id],
|
||||
size: len(left) + len(right),
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
for i := range len(runes) - 1 {
|
||||
if pair := pairwise(i, i+1); pair != nil {
|
||||
heap.Push(q, pair)
|
||||
}
|
||||
}
|
||||
|
||||
for q.Len() > 0 {
|
||||
pair := heap.Pop(q).(*candidate)
|
||||
left, right := merges[pair.a], merges[pair.b]
|
||||
|
||||
if string(left.runes) == "" || string(right.runes) == "" || len(string(left.runes))+len(string(right.runes)) != pair.size {
|
||||
continue
|
||||
}
|
||||
|
||||
merges[pair.a].runes = append(left.runes, right.runes...)
|
||||
merges[pair.b].runes = nil
|
||||
merges[pair.a].n = right.n
|
||||
if right.n < len(merges) {
|
||||
merges[right.n].p = pair.a
|
||||
}
|
||||
|
||||
if pair := pairwise(merges[pair.a].p, pair.a); pair != nil {
|
||||
heap.Push(q, pair)
|
||||
}
|
||||
|
||||
if pair := pairwise(pair.a, merges[pair.a].n); pair != nil {
|
||||
heap.Push(q, pair)
|
||||
}
|
||||
}
|
||||
|
||||
for _, merge := range merges {
|
||||
if token := string(merge.runes); token != "" {
|
||||
id := spm.vocab.Encode(token)
|
||||
|
||||
if id >= 0 {
|
||||
ids = append(ids, id)
|
||||
continue
|
||||
}
|
||||
|
||||
// Fallback to byte tokenization
|
||||
var result []int32
|
||||
for _, b := range []byte(token) {
|
||||
byteToken := fmt.Sprintf("<0x%02X>", b)
|
||||
unknownID := spm.vocab.Encode(byteToken)
|
||||
if unknownID >= 0 {
|
||||
result = append(result, unknownID)
|
||||
} else {
|
||||
slog.Debug("unknown byte token", "byte", b, "token", byteToken)
|
||||
}
|
||||
}
|
||||
|
||||
ids = append(ids, result...)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if addSpecial {
|
||||
ids = spm.vocab.addSpecials(ids)
|
||||
}
|
||||
|
||||
logutil.Trace("encoded", "string", s, "ids", ids)
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
type candidate struct {
|
||||
a, b int
|
||||
score float32
|
||||
size int
|
||||
}
|
||||
|
||||
type queue []*candidate
|
||||
|
||||
func (q queue) Len() int { return len(q) }
|
||||
|
||||
func (q queue) Less(i, j int) bool {
|
||||
return (q[i].score > q[j].score) || (q[i].score == q[j].score && q[i].a < q[j].a)
|
||||
}
|
||||
|
||||
func (q queue) Swap(i, j int) { q[i], q[j] = q[j], q[i] }
|
||||
|
||||
func (q *queue) Push(x interface{}) {
|
||||
item := x.(*candidate)
|
||||
*q = append(*q, item)
|
||||
}
|
||||
|
||||
func (q *queue) Pop() interface{} {
|
||||
old := *q
|
||||
n := len(old)
|
||||
item := old[n-1]
|
||||
*q = old[0 : n-1]
|
||||
return item
|
||||
}
|
||||
|
||||
func (spm SentencePiece) Decode(ids []int32) (string, error) {
|
||||
var sb strings.Builder
|
||||
for _, id := range ids {
|
||||
data := spm.vocab.Decode(id)
|
||||
data = strings.ReplaceAll(data, spmWhitespaceSep, " ")
|
||||
|
||||
// For tokenizer that use byte tokens like "<0xEA>"
|
||||
// convert them to the partial unicode character
|
||||
// so they are buffered correctly by the runner instead
|
||||
// of being sent back to the api as "<0xEA>"
|
||||
if len(data) == 6 && strings.HasPrefix(data, "<0x") && strings.HasSuffix(data, ">") {
|
||||
byteVal, err := strconv.ParseUint(data[1:5], 0, 8)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("failed to parse hex byte: %v", err)
|
||||
}
|
||||
|
||||
if err := sb.WriteByte(byte(byteVal)); err != nil {
|
||||
return "", err
|
||||
}
|
||||
} else {
|
||||
if _, err := sb.WriteString(data); err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
logutil.Trace("decoded", "ids", ids, "string", sb.String())
|
||||
return sb.String(), nil
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,333 +0,0 @@
|
||||
// Copyright 2016 Google Inc.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.!
|
||||
|
||||
syntax = "proto2";
|
||||
|
||||
// TODO(taku): Needs to use LITE RUNTIME in OSS release.
|
||||
option optimize_for = LITE_RUNTIME;
|
||||
option go_package = "./sentencepiece";
|
||||
|
||||
package sentencepiece;
|
||||
|
||||
// TrainerSpec encodes a various parameters for SentencePiece training.
|
||||
// Next id: 55
|
||||
message TrainerSpec {
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// General parameters
|
||||
//
|
||||
// Input corpus files.
|
||||
// Trainer accepts the following two formats:
|
||||
// A) Monolingual: plain text, one sentence per line.
|
||||
// B) Bilingual: TSV, source sentence <tab> target sentence
|
||||
// When bilingual data is passed, shared vocabulary model is built.
|
||||
// Note that the input file must be raw corpus, not a preprocessed corpus.
|
||||
// Trainer only loads the first `input_sentence_size` sentences specified
|
||||
// with this parameter.
|
||||
repeated string input = 1;
|
||||
|
||||
// Input corpus format:
|
||||
// "text": one-sentence-per-line text format (default)
|
||||
// "tsv": sentence <tab> freq
|
||||
optional string input_format = 7;
|
||||
|
||||
// Output model file prefix.
|
||||
// <model_prefix>.model and <model_prefix>.vocab are generated.
|
||||
optional string model_prefix = 2;
|
||||
|
||||
// Model type. only have UNIGRAM now.
|
||||
enum ModelType {
|
||||
UNIGRAM = 1; // Unigram language model with dynamic algorithm
|
||||
BPE = 2; // Byte Pair Encoding
|
||||
WORD = 3; // Delimitered by whitespace.
|
||||
CHAR = 4; // tokenizes into character sequence
|
||||
}
|
||||
optional ModelType model_type = 3 [default = UNIGRAM];
|
||||
|
||||
// Vocabulary size. 8k is the default size.
|
||||
optional int32 vocab_size = 4 [default = 8000];
|
||||
|
||||
// List of the languages this model can accept.
|
||||
// Since the model is language-agnostic, this field is used as a reference.
|
||||
repeated string accept_language = 5;
|
||||
|
||||
// Size of self-test samples, which are encoded in the model file.
|
||||
optional int32 self_test_sample_size = 6 [default = 0];
|
||||
|
||||
// Whether to use DP version of sentencepiece. Use it with TSV input format
|
||||
// (requires precomputed word tab counts to work).
|
||||
optional bool enable_differential_privacy = 50 [default = false];
|
||||
// Set these parameters if you need DP version of sentencepiece.
|
||||
// std of noise to add.
|
||||
optional float differential_privacy_noise_level = 51 [default = 0.0];
|
||||
// Clipping threshold to apply after adding noise. All the words with
|
||||
// frequency less than this value are dropped.
|
||||
optional uint64 differential_privacy_clipping_threshold = 52 [default = 0];
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// Training parameters.
|
||||
//
|
||||
// Uses characters which cover the corpus with the ratio of `chars_coverage`.
|
||||
// This parameter determines the set of basic Alphabet of sentence piece.
|
||||
// 1.0 - `chars_coverage` characters are treated as UNK.
|
||||
// See also required_chars field.
|
||||
optional float character_coverage = 10 [default = 0.9995];
|
||||
|
||||
// Maximum size of sentences the trainer loads from `input` parameter.
|
||||
// Trainer simply loads the `input` files in sequence.
|
||||
// It is better to shuffle the input corpus randomly.
|
||||
optional uint64 input_sentence_size = 11 [default = 0];
|
||||
optional bool shuffle_input_sentence = 19 [default = true];
|
||||
|
||||
// Maximum size of sentences to make seed sentence pieces.
|
||||
// Extended suffix array is constructed to extract frequent
|
||||
// sub-strings from the corpus. This uses 20N working space,
|
||||
// where N is the size of corpus.
|
||||
optional int32 mining_sentence_size = 12 [deprecated = true];
|
||||
|
||||
// Maximum size of sentences to train sentence pieces.
|
||||
optional int32 training_sentence_size = 13 [deprecated = true];
|
||||
|
||||
// The size of seed sentencepieces.
|
||||
// `seed_sentencepiece_size` must be larger than `vocab_size`.
|
||||
optional int32 seed_sentencepiece_size = 14 [default = 1000000];
|
||||
|
||||
// In every EM sub-iterations, keeps top
|
||||
// `shrinking_factor` * `current sentencepieces size` with respect to
|
||||
// the loss of the sentence piece. This value should be smaller than 1.0.
|
||||
optional float shrinking_factor = 15 [default = 0.75];
|
||||
|
||||
// The maximum sentence length in byte. The sentences with the length
|
||||
// larger than `max_sentence_length` is simply ignored.
|
||||
// Longer input tends to bring the following risks:
|
||||
// * Overflow during EM training (unigram language model only)
|
||||
// * Performance drop because of O(n log n) cost in BPE.
|
||||
optional int32 max_sentence_length = 18 [default = 4192];
|
||||
|
||||
// Number of threads in the training.
|
||||
optional int32 num_threads = 16 [default = 16];
|
||||
|
||||
// Number of EM sub iterations.
|
||||
optional int32 num_sub_iterations = 17 [default = 2];
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// SentencePiece parameters which control the shapes of sentence piece.
|
||||
//
|
||||
// Maximum length of sentencepiece.
|
||||
optional int32 max_sentencepiece_length = 20 [default = 16];
|
||||
|
||||
// Uses Unicode script to split sentence pieces.
|
||||
// When `split_by_unicode_script` is true, we do not allow sentence piece to
|
||||
// include multiple Unicode scripts, e.g. "F1" is not a valid piece.
|
||||
// Exception: CJ characters (Hiragana/Katakana/Han) are all handled
|
||||
// as one script type, since Japanese word can consist of multiple scripts.
|
||||
// This exception is always applied regardless of the accept-language
|
||||
// parameter.
|
||||
optional bool split_by_unicode_script = 21 [default = true];
|
||||
|
||||
// When `split_by_number` is true, put a boundary between number and
|
||||
// non-number transition. If we want to treat "F1" is one token, set this flag
|
||||
// to be false.
|
||||
optional bool split_by_number = 23 [default = true];
|
||||
|
||||
// Use a white space to split sentence pieces.
|
||||
// When `split_by_whitespace` is false, we may have the piece containing
|
||||
// a white space in the middle. e.g., "in_the".
|
||||
optional bool split_by_whitespace = 22 [default = true];
|
||||
|
||||
// Adds whitespace symbol (_) as a suffix instead of prefix. e.g., _hello =>
|
||||
// hello_. When `treat_whitespace_as_suffix` is true,
|
||||
// NormalizerSpec::add_dummy_prefix will add the dummy whitespace to the end
|
||||
// of sentence.
|
||||
optional bool treat_whitespace_as_suffix = 24 [default = false];
|
||||
|
||||
// Allows pieces that only contain whitespaces instead of appearing only as
|
||||
// prefix or suffix of other pieces.
|
||||
optional bool allow_whitespace_only_pieces = 26 [default = false];
|
||||
|
||||
// Split all digits (0-9) into separate pieces.
|
||||
optional bool split_digits = 25 [default = false];
|
||||
|
||||
// Defines the pre-tokenization delimiter.
|
||||
// When specified, no pieces crossing this delimiter is not included
|
||||
// in the vocab. Then the delimiter string is virtually ignored
|
||||
// during the training. This field can allows constraints on the vocabulary
|
||||
// selection. Note that this field is available on unigram mode.
|
||||
optional string pretokenization_delimiter = 53 [ default = ""];
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// Vocabulary management
|
||||
//
|
||||
// Defines control symbols used as an indicator to
|
||||
// change the behavior of the decoder. <s> and </s> are pre-defined.
|
||||
// We can use this field to encode various meta information,
|
||||
// including language indicator in multilingual model.
|
||||
// These symbols are not visible to users, but visible to
|
||||
// the decoder. Note that when the input sentence contains control symbols,
|
||||
// they are not treated as one token, but segmented into normal pieces.
|
||||
// Control symbols must be inserted independently from the segmentation.
|
||||
repeated string control_symbols = 30;
|
||||
|
||||
// Defines user defined symbols.
|
||||
// These symbols are added with extremely high score
|
||||
// so they are always treated as one unique symbol in any context.
|
||||
// Typical usage of user_defined_symbols is placeholder for named entities.
|
||||
repeated string user_defined_symbols = 31;
|
||||
|
||||
// Defines required characters. Each UTF8 character in this string is included
|
||||
// in the character set regardless of character_coverage value. Unlike
|
||||
// user_defined_symbols, these characters have scores based on the frequency
|
||||
// on input sentences, and the model can form subwords using characters
|
||||
// in this field.
|
||||
optional string required_chars = 36;
|
||||
|
||||
// Decomposes unknown pieces into UTF-8 bytes.
|
||||
optional bool byte_fallback = 35 [default = false];
|
||||
|
||||
// When creating the vocabulary file, defines whether or not to additionally
|
||||
// output the score for each piece.
|
||||
optional bool vocabulary_output_piece_score = 32 [default = true];
|
||||
|
||||
// `vocab_size` is treated as hard limit. Crash if
|
||||
// the model can not produce the vocab of size `vocab_size`,
|
||||
// When `hard_vocab_limit` is false, vocab_size is treated
|
||||
// as soft limit. Note that when model_type=char,
|
||||
// always assumes hard_vocab_limit = false.
|
||||
optional bool hard_vocab_limit = 33 [default = true];
|
||||
|
||||
// use all symbols for vocab extraction. This flag is valid
|
||||
// if model type is either CHAR or WORD
|
||||
optional bool use_all_vocab = 34 [default = false];
|
||||
|
||||
///////////////////////////////////////////////////////////////////
|
||||
// Reserved special meta tokens.
|
||||
// * -1 is not used.
|
||||
// * unk_id must not be -1.
|
||||
// Id must start with 0 and be contiguous.
|
||||
optional int32 unk_id = 40 [default = 0]; // <unk>
|
||||
optional int32 bos_id = 41 [default = 1]; // <s>
|
||||
optional int32 eos_id = 42 [default = 2]; // </s>
|
||||
optional int32 pad_id = 43 [default = -1]; // <pad> (padding)
|
||||
optional string unk_piece = 45 [default = "<unk>"];
|
||||
optional string bos_piece = 46 [default = "<s>"];
|
||||
optional string eos_piece = 47 [default = "</s>"];
|
||||
optional string pad_piece = 48 [default = "<pad>"];
|
||||
|
||||
// Encodes <unk> into U+2047 (DOUBLE QUESTION MARK),
|
||||
// since this character can be useful both for user and
|
||||
// developer. We can easily figure out that <unk> is emitted.
|
||||
optional string unk_surface = 44 [default = " \xE2\x81\x87 "];
|
||||
|
||||
// Increase bit depth to allow unigram model training on large
|
||||
// (>10M sentences) corpora. A Side-effect of enabling this flag
|
||||
// is increased memory usage.
|
||||
optional bool train_extremely_large_corpus = 49 [default = false];
|
||||
|
||||
// Path to a seed sentencepieces file, with one tab-separated
|
||||
// seed sentencepiece <tab> frequency per line.
|
||||
optional string seed_sentencepieces_file = 54 [default = ""];
|
||||
|
||||
// Customized extensions: the range of field numbers
|
||||
// are open to third-party extensions.
|
||||
extensions 200 to max;
|
||||
}
|
||||
|
||||
// NormalizerSpec encodes a various parameters for string normalizaiton
|
||||
message NormalizerSpec {
|
||||
// name of normalization rule.
|
||||
optional string name = 1;
|
||||
|
||||
// Pre-compiled normalization rule created by
|
||||
// Builder::GetPrecompiledCharsMap() or Builder::CompileCharsMap() method.
|
||||
// Usually this field is set by Builder::GetNormalizerSpec() method.
|
||||
optional bytes precompiled_charsmap = 2;
|
||||
|
||||
// Adds dummy whitespace at the beginning of text in order to
|
||||
// treat "world" in "world" and "hello world" in the same way.
|
||||
optional bool add_dummy_prefix = 3 [default = true];
|
||||
|
||||
// Removes leading, trailing, and duplicate internal whitespace.
|
||||
optional bool remove_extra_whitespaces = 4 [default = true];
|
||||
|
||||
// Replaces whitespace with meta symbol.
|
||||
// This field must be true to train sentence piece model.
|
||||
optional bool escape_whitespaces = 5 [default = true];
|
||||
|
||||
// Custom normalization rule file in TSV format.
|
||||
// https://github.com/google/sentencepiece/blob/master/doc/normalization.md
|
||||
// This field is only used in SentencePieceTrainer::Train() method, which
|
||||
// compiles the rule into the binary rule stored in `precompiled_charsmap`.
|
||||
optional string normalization_rule_tsv = 6;
|
||||
|
||||
// Customized extensions: the range of field numbers
|
||||
// are open to third-party extensions.
|
||||
extensions 200 to max;
|
||||
}
|
||||
|
||||
// Proto to store samples for self-testing.
|
||||
message SelfTestData {
|
||||
message Sample {
|
||||
optional string input = 1;
|
||||
optional string expected = 2;
|
||||
}
|
||||
repeated Sample samples = 1;
|
||||
|
||||
// Customized extensions: the range of field numbers
|
||||
// are open to third-party extensions.
|
||||
extensions 200 to max;
|
||||
}
|
||||
|
||||
// ModelProto stores model parameters.
|
||||
// SentencePieceProcessor is supposed to be self-contained.
|
||||
// All settings/parameters which may change the behavior must be encoded
|
||||
// in ModelProto.
|
||||
message ModelProto {
|
||||
message SentencePiece {
|
||||
enum Type {
|
||||
NORMAL = 1; // normal symbol
|
||||
UNKNOWN = 2; // unknown symbol. only <unk> for now.
|
||||
CONTROL = 3; // control symbols. </s>, <s>, <2ja> etc.
|
||||
USER_DEFINED = 4; // user defined symbols.
|
||||
// Typical usage of USER_DEFINED symbol
|
||||
// is placeholder.
|
||||
BYTE = 6; // byte symbols. Used when `byte_fallback` is true.
|
||||
UNUSED = 5; // this piece is not used.
|
||||
}
|
||||
optional string piece = 1; // piece must not be empty.
|
||||
optional float score = 2;
|
||||
optional Type type = 3 [default = NORMAL];
|
||||
|
||||
// Customized extensions: the range of field numbers
|
||||
// are open to third-party extensions.
|
||||
extensions 200 to max;
|
||||
}
|
||||
|
||||
// Sentence pieces with scores.
|
||||
repeated SentencePiece pieces = 1;
|
||||
|
||||
// Spec used to generate this model file.
|
||||
optional TrainerSpec trainer_spec = 2;
|
||||
|
||||
// Spec for text normalization.
|
||||
optional NormalizerSpec normalizer_spec = 3;
|
||||
|
||||
// Stores sample input and its expected segmentation to verify the model.
|
||||
optional SelfTestData self_test_data = 4;
|
||||
|
||||
// Spec for text de-normalization.
|
||||
optional NormalizerSpec denormalizer_spec = 5;
|
||||
|
||||
// Customized extensions: the range of field numbers
|
||||
// are open to third-party extensions.
|
||||
extensions 200 to max;
|
||||
}
|
||||
@@ -1,172 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"slices"
|
||||
"testing"
|
||||
|
||||
"google.golang.org/protobuf/proto"
|
||||
|
||||
"github.com/ollama/ollama/tokenizer/sentencepiece"
|
||||
)
|
||||
|
||||
func loadSentencePieceVocab(t *testing.T) SentencePiece {
|
||||
t.Helper()
|
||||
|
||||
bts, err := os.ReadFile(filepath.FromSlash("testdata/gemma2/tokenizer.model"))
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var spm sentencepiece.ModelProto
|
||||
if err := proto.Unmarshal(bts, &spm); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
var v Vocabulary
|
||||
|
||||
for _, piece := range spm.GetPieces() {
|
||||
v.Values = append(v.Values, piece.GetPiece())
|
||||
v.Scores = append(v.Scores, piece.GetScore())
|
||||
switch t := piece.GetType(); t {
|
||||
case sentencepiece.ModelProto_SentencePiece_UNKNOWN,
|
||||
sentencepiece.ModelProto_SentencePiece_CONTROL,
|
||||
sentencepiece.ModelProto_SentencePiece_UNUSED,
|
||||
sentencepiece.ModelProto_SentencePiece_BYTE:
|
||||
v.Types = append(v.Types, int32(t))
|
||||
default:
|
||||
tt := int32(sentencepiece.ModelProto_SentencePiece_NORMAL)
|
||||
// todo parse the special tokens file
|
||||
// - this will roundtrip correctly but the <start_of_turn> and
|
||||
// <end_of_turn> tokens aren't processed
|
||||
v.Types = append(v.Types, tt)
|
||||
}
|
||||
}
|
||||
|
||||
return NewSentencePiece(&v)
|
||||
}
|
||||
|
||||
func TestSentencePieceEncode(t *testing.T) {
|
||||
logger := slog.New(slog.NewTextHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug}))
|
||||
slog.SetDefault(logger)
|
||||
|
||||
tokenizer := loadSentencePieceVocab(t)
|
||||
|
||||
t.Run("basic roundtrip", func(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
cases := []string{
|
||||
"hello",
|
||||
"hello ",
|
||||
"hello ",
|
||||
" hello",
|
||||
" hello ",
|
||||
" hello ",
|
||||
"hello world",
|
||||
"请考试我的软件！12345",
|
||||
"你好",
|
||||
"Hello 你好 world!",
|
||||
"Special characters: !@#$%^&*()_+-=[]{}|;':\",./<>?",
|
||||
"Multilingual: 你好 こんにちは Привет Hola مرحبا",
|
||||
"Numbers and symbols: 123456789 +- */",
|
||||
"Special tokens: <bos> text <eos>",
|
||||
"Code snippets: func main() { fmt.Println(\"Hello World\") }",
|
||||
"Long text: " + "Lorem ipsum dolor sit amet, consectetur adipiscing elit. " +
|
||||
"Sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. " +
|
||||
"Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris.",
|
||||
}
|
||||
|
||||
for _, want := range cases {
|
||||
ids, err := tokenizer.Encode(want, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if got, err := tokenizer.Decode(ids); err != nil {
|
||||
t.Fatal(err)
|
||||
} else if got != want {
|
||||
t.Errorf("got %q, want %q [%#v]", got, want, ids)
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("special tokens", func(t *testing.T) {
|
||||
type candidate struct {
|
||||
token string
|
||||
ids []int32
|
||||
}
|
||||
|
||||
cases := []candidate{
|
||||
{"<bos>", []int32{2}},
|
||||
{"<eos>", []int32{1}},
|
||||
}
|
||||
|
||||
for _, want := range cases {
|
||||
ids, err := tokenizer.Encode(want.token, true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if !slices.Equal(ids, want.ids) {
|
||||
t.Errorf("got %#v, want %#v", ids, want.ids)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestSentencePieceDecodeByteTokens(t *testing.T) {
|
||||
vocab := &Vocabulary{
|
||||
Values: []string{
|
||||
"normal",
|
||||
"<0xEA>",
|
||||
"<0x41>",
|
||||
"<0xC3>",
|
||||
"<0xA3>",
|
||||
},
|
||||
Types: []int32{
|
||||
TOKEN_TYPE_NORMAL,
|
||||
TOKEN_TYPE_BYTE,
|
||||
TOKEN_TYPE_BYTE,
|
||||
TOKEN_TYPE_BYTE,
|
||||
TOKEN_TYPE_BYTE,
|
||||
},
|
||||
Scores: []float32{0, 0, 0, 0, 0},
|
||||
}
|
||||
|
||||
spm := NewSentencePiece(vocab)
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
ids []int32
|
||||
expected string
|
||||
}{
|
||||
{
|
||||
name: "single byte token",
|
||||
ids: []int32{1},
|
||||
expected: "\xea",
|
||||
},
|
||||
{
|
||||
name: "ASCII byte token",
|
||||
ids: []int32{2},
|
||||
expected: "A",
|
||||
},
|
||||
{
|
||||
name: "multiple byte tokens forming UTF-8 character",
|
||||
ids: []int32{3, 4},
|
||||
expected: "ã",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result, err := spm.Decode(tt.ids)
|
||||
if err != nil {
|
||||
t.Errorf("failed to decode token IDs %v: %v", tt.ids, err)
|
||||
}
|
||||
if result != tt.expected {
|
||||
t.Errorf("got %q, want %q", result, tt.expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
BIN
Binary file not shown.
Vendored
-63845
File diff suppressed because it is too large
Load Diff
@@ -1,15 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
const (
|
||||
TOKEN_TYPE_NORMAL = iota + 1
|
||||
TOKEN_TYPE_UNKNOWN
|
||||
TOKEN_TYPE_CONTROL
|
||||
TOKEN_TYPE_USER_DEFINED
|
||||
TOKEN_TYPE_UNUSED
|
||||
TOKEN_TYPE_BYTE
|
||||
)
|
||||
|
||||
type Tokenizer interface {
|
||||
Encode(s string, addSpecial bool) ([]int32, error)
|
||||
Decode([]int32) (string, error)
|
||||
}
|
||||
@@ -1,94 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"slices"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type Vocabulary struct {
|
||||
Values []string
|
||||
Types []int32
|
||||
Scores []float32
|
||||
Merges []string
|
||||
|
||||
BOS, EOS []int32
|
||||
AddBOS, AddEOS bool
|
||||
|
||||
specialOnce sync.Once
|
||||
special []string
|
||||
|
||||
valuesOnce sync.Once
|
||||
values map[string]int32
|
||||
|
||||
mergeOnce sync.Once
|
||||
merge map[string]int32
|
||||
}
|
||||
|
||||
func (v *Vocabulary) addSpecials(ids []int32) []int32 {
|
||||
if v.AddBOS && len(v.BOS) > 0 {
|
||||
if len(ids) > 0 && slices.Contains(v.BOS, ids[0]) {
|
||||
slog.Warn("adding bos token to prompt which already has it", "id", v.BOS)
|
||||
}
|
||||
|
||||
slog.Debug("adding bos token to prompt", "id", v.BOS[0])
|
||||
ids = append([]int32{v.BOS[0]}, ids...)
|
||||
}
|
||||
|
||||
if v.AddEOS && len(v.EOS) > 0 {
|
||||
if len(ids) > 0 && slices.Contains(v.BOS, ids[len(ids)-1]) {
|
||||
slog.Warn("adding eos token to prompt which already has it", "id", v.EOS)
|
||||
}
|
||||
|
||||
slog.Debug("adding eos token to prompt", "id", v.EOS[0])
|
||||
ids = append(ids, v.EOS[0])
|
||||
}
|
||||
|
||||
return ids
|
||||
}
|
||||
|
||||
func (v *Vocabulary) Encode(s string) int32 {
|
||||
v.valuesOnce.Do(func() {
|
||||
v.values = make(map[string]int32, len(v.Values))
|
||||
for i, value := range v.Values {
|
||||
v.values[value] = int32(i)
|
||||
}
|
||||
})
|
||||
|
||||
if id, ok := v.values[s]; ok {
|
||||
return id
|
||||
}
|
||||
|
||||
return -1
|
||||
}
|
||||
|
||||
func (v *Vocabulary) Decode(id int32) string {
|
||||
return v.Values[id]
|
||||
}
|
||||
|
||||
func (v *Vocabulary) SpecialVocabulary() []string {
|
||||
v.specialOnce.Do(func() {
|
||||
for i := range v.Values {
|
||||
if v.Types[i] == TOKEN_TYPE_CONTROL || v.Types[i] == TOKEN_TYPE_USER_DEFINED {
|
||||
v.special = append(v.special, v.Values[i])
|
||||
}
|
||||
}
|
||||
})
|
||||
|
||||
return v.special
|
||||
}
|
||||
|
||||
func (v *Vocabulary) Merge(left, right string) int {
|
||||
v.mergeOnce.Do(func() {
|
||||
v.merge = make(map[string]int32, len(v.Merges))
|
||||
for i, merge := range v.Merges {
|
||||
v.merge[merge] = int32(i)
|
||||
}
|
||||
})
|
||||
|
||||
if id, ok := v.merge[left+" "+right]; ok {
|
||||
return int(id)
|
||||
}
|
||||
|
||||
return -1
|
||||
}
|
||||
@@ -1,107 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
)
|
||||
|
||||
func TestSpecialVocabulary(t *testing.T) {
|
||||
vocab := &Vocabulary{
|
||||
Values: []string{"<|startoftext|>", "<|endoftext|>", "<|tool_call_start|>", "<|tool_call_end|>", "hi"},
|
||||
Types: []int32{TOKEN_TYPE_CONTROL, TOKEN_TYPE_CONTROL, TOKEN_TYPE_USER_DEFINED, TOKEN_TYPE_USER_DEFINED, TOKEN_TYPE_NORMAL},
|
||||
}
|
||||
|
||||
specialVocab := vocab.SpecialVocabulary()
|
||||
|
||||
if len(specialVocab) != 4 {
|
||||
t.Errorf("expected 4 special tokens, got %d", len(specialVocab))
|
||||
}
|
||||
}
|
||||
|
||||
func TestAddSpecialVocabulary(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
vocab *Vocabulary
|
||||
input []int32
|
||||
want []int32
|
||||
}{
|
||||
{
|
||||
name: "add bos",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: true,
|
||||
AddEOS: false,
|
||||
},
|
||||
input: []int32{2, 3, 4},
|
||||
want: []int32{0, 2, 3, 4},
|
||||
},
|
||||
{
|
||||
// TODO(mxyng): this is to match previous behaviour
|
||||
name: "add bos when already present",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: true,
|
||||
AddEOS: false,
|
||||
},
|
||||
input: []int32{0, 2, 3, 4},
|
||||
want: []int32{0, 0, 2, 3, 4},
|
||||
},
|
||||
{
|
||||
name: "add eos",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: false,
|
||||
AddEOS: true,
|
||||
},
|
||||
input: []int32{2, 3, 4},
|
||||
want: []int32{2, 3, 4, 1},
|
||||
},
|
||||
{
|
||||
// TODO(mxyng): this is to match previous behaviour
|
||||
name: "add eos when already present",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: false,
|
||||
AddEOS: true,
|
||||
},
|
||||
input: []int32{2, 3, 4, 1},
|
||||
want: []int32{2, 3, 4, 1, 1},
|
||||
},
|
||||
{
|
||||
name: "add both",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: true,
|
||||
AddEOS: true,
|
||||
},
|
||||
input: []int32{2, 3, 4},
|
||||
want: []int32{0, 2, 3, 4, 1},
|
||||
},
|
||||
{
|
||||
name: "add bos to empty inputs",
|
||||
vocab: &Vocabulary{
|
||||
BOS: []int32{0},
|
||||
EOS: []int32{1},
|
||||
AddBOS: true,
|
||||
AddEOS: false,
|
||||
},
|
||||
input: []int32{},
|
||||
want: []int32{0},
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range cases {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
got := tt.vocab.addSpecials(tt.input)
|
||||
if diff := cmp.Diff(tt.want, got); diff != "" {
|
||||
t.Errorf("no match (-want +got):\n%s", diff)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -1,161 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"iter"
|
||||
"strings"
|
||||
"unicode"
|
||||
|
||||
"github.com/ollama/ollama/logutil"
|
||||
)
|
||||
|
||||
type WordPiece struct {
|
||||
vocab *Vocabulary
|
||||
lowercase bool
|
||||
}
|
||||
|
||||
// ggmlPrefix is the prefix used by GGML vocabularies to indicate word boundaries.
|
||||
// this differs from original word piece which uses "##" to indicate subwords.
|
||||
const ggmlPrefix = "▁"
|
||||
|
||||
var wordPieceReplacer = strings.NewReplacer(
|
||||
" .", ".",
|
||||
" ?", "?",
|
||||
" !", "!",
|
||||
" ,", ",",
|
||||
" ' ", "'",
|
||||
" n't", "n't",
|
||||
" 'm", "'m",
|
||||
" do not", " don't",
|
||||
" 's", "'s",
|
||||
" 've", "'ve",
|
||||
" 're", "'re",
|
||||
)
|
||||
|
||||
// Decode implements Tokenizer.
|
||||
func (wpm WordPiece) Decode(ids []int32) (string, error) {
|
||||
var sb strings.Builder
|
||||
for i, id := range ids {
|
||||
if id < 0 || int(id) >= len(wpm.vocab.Values) {
|
||||
return "", fmt.Errorf("invalid token id: %d", id)
|
||||
}
|
||||
|
||||
var separator string
|
||||
piece := wpm.vocab.Values[id]
|
||||
if i > 0 &&
|
||||
(strings.HasPrefix(piece, ggmlPrefix) ||
|
||||
(strings.HasPrefix(piece, "[") && strings.HasSuffix(piece, "]"))) {
|
||||
separator = " "
|
||||
}
|
||||
|
||||
sb.WriteString(wordPieceReplacer.Replace(separator + strings.TrimPrefix(piece, ggmlPrefix)))
|
||||
}
|
||||
|
||||
return sb.String(), nil
|
||||
}
|
||||
|
||||
// words splits a string into words, treating CJK characters as separate words.
|
||||
// TODO: this is specifically for BERT and may need to be adjusted or refactored for other models.
|
||||
func (wpm WordPiece) words(s string) iter.Seq[string] {
|
||||
return func(yield func(string) bool) {
|
||||
runes := make([]rune, 0, len(s)*3)
|
||||
for _, r := range s {
|
||||
switch {
|
||||
case r >= 0x4E00 && r <= 0x9FFF,
|
||||
r >= 0x3400 && r <= 0x4DBF,
|
||||
r >= 0x20000 && r <= 0x2A6DF,
|
||||
r >= 0x2A700 && r <= 0x2B73F,
|
||||
r >= 0x2B740 && r <= 0x2B81F,
|
||||
r >= 0x2B820 && r <= 0x2CEAF,
|
||||
r >= 0xF900 && r <= 0xFAFF,
|
||||
r >= 0x2F800 && r <= 0x2FA1F:
|
||||
runes = append(runes, ' ', r, ' ')
|
||||
default:
|
||||
runes = append(runes, r)
|
||||
}
|
||||
}
|
||||
|
||||
for w := range strings.FieldsFuncSeq(string(runes), unicode.IsSpace) {
|
||||
// split on but keep punctuation
|
||||
var start int
|
||||
for start < len(w) {
|
||||
end := strings.IndexFunc(w[start:], unicode.IsPunct)
|
||||
if end < 0 {
|
||||
end = len(w) - start
|
||||
} else if end == 0 {
|
||||
end = 1
|
||||
}
|
||||
|
||||
if !yield(w[start : start+end]) {
|
||||
return
|
||||
}
|
||||
|
||||
start += end
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Encode implements Tokenizer.
|
||||
func (wpm WordPiece) Encode(s string, addSpecial bool) ([]int32, error) {
|
||||
var ids []int32
|
||||
|
||||
// TODO: use [UNK] from config
|
||||
unk := wpm.vocab.Encode("[UNK]")
|
||||
for word := range wpm.words(s) {
|
||||
var start int
|
||||
var pieces []int32
|
||||
for start < len(word) {
|
||||
end := len(word)
|
||||
|
||||
var piece int32
|
||||
for start < end {
|
||||
subword := word[start:end]
|
||||
if start == 0 {
|
||||
subword = ggmlPrefix + subword
|
||||
}
|
||||
|
||||
if wpm.lowercase {
|
||||
subword = strings.ToLower(subword)
|
||||
}
|
||||
piece = wpm.vocab.Encode(subword)
|
||||
if piece >= 0 {
|
||||
break
|
||||
}
|
||||
|
||||
end--
|
||||
}
|
||||
|
||||
if piece < 0 {
|
||||
// Unknown token
|
||||
pieces = pieces[:0]
|
||||
break
|
||||
}
|
||||
|
||||
pieces = append(pieces, piece)
|
||||
start = end
|
||||
}
|
||||
|
||||
if len(pieces) > 0 {
|
||||
ids = append(ids, pieces...)
|
||||
} else {
|
||||
ids = append(ids, unk)
|
||||
}
|
||||
}
|
||||
|
||||
if addSpecial {
|
||||
ids = wpm.vocab.addSpecials(ids)
|
||||
}
|
||||
|
||||
logutil.Trace("encoded", "string", s, "ids", ids)
|
||||
return ids, nil
|
||||
}
|
||||
|
||||
var _ Tokenizer = (*WordPiece)(nil)
|
||||
|
||||
func NewWordPiece(vocab *Vocabulary, lowercase bool) WordPiece {
|
||||
return WordPiece{
|
||||
vocab: vocab,
|
||||
lowercase: lowercase,
|
||||
}
|
||||
}
|
||||
@@ -1,53 +0,0 @@
|
||||
package tokenizer
|
||||
|
||||
import (
|
||||
"slices"
|
||||
"testing"
|
||||
|
||||
"github.com/google/go-cmp/cmp"
|
||||
)
|
||||
|
||||
func TestWordPiece(t *testing.T) {
|
||||
wpm := NewWordPiece(
|
||||
&Vocabulary{
|
||||
Values: []string{"[UNK]", "[CLS]", "[SEP]", "▁hello", "▁world", "s", "▁!", "▁@", "▁#"},
|
||||
AddBOS: true,
|
||||
AddEOS: true,
|
||||
BOS: []int32{1},
|
||||
EOS: []int32{2},
|
||||
},
|
||||
true, // lowercase
|
||||
)
|
||||
|
||||
ids, err := wpm.Encode("Hello world!", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff([]int32{1, 3, 4, 6, 2}, ids); diff != "" {
|
||||
t.Errorf("unexpected ids (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
words, err := wpm.Decode(ids)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if diff := cmp.Diff("[CLS] hello world! [SEP]", words); diff != "" {
|
||||
t.Errorf("unexpected words (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
|
||||
func TestWordPieceWords(t *testing.T) {
|
||||
var wpm WordPiece
|
||||
|
||||
basic := slices.Collect(wpm.words("Hey friend! How are you?!?"))
|
||||
if diff := cmp.Diff([]string{"Hey", "friend", "!", "How", "are", "you", "?", "!", "?"}, basic); diff != "" {
|
||||
t.Errorf("unexpected words (-want +got):\n%s", diff)
|
||||
}
|
||||
|
||||
chinese := slices.Collect(wpm.words("野口里佳 Noguchi Rika"))
|
||||
if diff := cmp.Diff([]string{"野", "口", "里", "佳", "Noguchi", "Rika"}, chinese); diff != "" {
|
||||
t.Errorf("unexpected words (-want +got):\n%s", diff)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user