embedder

package
v0.3.1 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Sep 1, 2026 License: MIT Imports: 22 Imported by: 0

README

embedder

The Embedder interface and its implementations. Stores and pipelines receive an Embedder by dependency injection — recall itself never calls an LLM API directly.

type Embedder interface {
    Embed(ctx context.Context, text string) ([]float32, error)
    EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)
    Dimension() int
}

Implementations

Embedder Constructor Notes
MockEmbedder NewMockEmbedder(dim) Deterministic hash-based vectors; for tests and offline examples.
OpenAIEmbedder NewOpenAIEmbedder(cfg) text-embedding-3-* models; Dimension validated against the model.
CohereEmbedder NewCohereEmbedder(cfg) embed-english-v3.0 / embed-multilingual-v3.0.
OllamaEmbedder NewOllamaEmbedder(cfg) Local Ollama server.
OnnxEmbedder NewOnnxEmbedder(cfg) / NewOnnxEmbedderFile(path, cfg) Local ONNX inference (see embedder/onnx), no CGO.
CachingEmbedder NewCachingEmbedder(inner, cache, ttl) Wraps any embedder with an LRU+TTL cache (cache package).

MultiModalEmbedder + MockMultiModal cover image+text embedding for the multimodal store. Helpers: CosineSimilarity, AutoDimension, and Hugging Face model resolution (LoadHFModel, ModelCache) for ONNX models.

API keys

Provider embedders take the key in their config struct; the config / app layers read it from the environment variable named by EmbedderConfig.APIKeyEnv. Keys are never stored in config files.

Documentation

Overview

Package embedder provides the interface for generating text embeddings. Users inject their preferred embedding implementation (OpenAI, local model, mock, etc.).

Index

Constants

View Source 
const (
	// ModelEmbedEnglishV3 is Cohere's English embedding model.
	ModelEmbedEnglishV3 = "embed-english-v3.0"

	// ModelEmbedMultilingualV3 is Cohere's multilingual embedding model.
	ModelEmbedMultilingualV3 = "embed-multilingual-v3.0"
)

Well-known Cohere embedding model names (both 1024 dimensions).

View Source 
const (
	// InputTypeSearchDocument marks text being embedded for indexing.
	InputTypeSearchDocument = "search_document"

	// InputTypeSearchQuery marks text being embedded as a search query.
	InputTypeSearchQuery = "search_query"

	// InputTypeClassification marks text being embedded for classification.
	InputTypeClassification = "classification"
)

Cohere input types (the "input_type" request field).

View Source 
const (
	// TruncationNone disables truncation (the default).
	TruncationNone = "NONE"

	// TruncationStart truncates the beginning of the input.
	TruncationStart = "START"

	// TruncationEnd truncates the end of the input.
	TruncationEnd = "END"
)

Cohere truncation strategies (the "truncate" request field).

View Source 
const (
	// ModelAllMiniLML6V2 is a fast, general-purpose sentence model (384 dims).
	ModelAllMiniLML6V2 = "all-minilm-l6-v2"

	// ModelNomicEmbedText is a high-quality general embedding model (768 dims).
	ModelNomicEmbedText = "nomic-embed-text"

	// ModelBGESmallENV15 is a compact BGE model (384 dims).
	ModelBGESmallENV15 = "bge-small-en-v1.5"
)

Well-known local embedding models available through Ollama, with their output dimensions.

View Source 
const (
	// ModelTextEmbedding3Small is the default OpenAI embedding model (1536 dims).
	ModelTextEmbedding3Small = "text-embedding-3-small"

	// ModelTextEmbedding3Large is the highest-quality OpenAI embedding model (3072 dims).
	ModelTextEmbedding3Large = "text-embedding-3-large"

	// ModelTextEmbeddingAda002 is the previous-generation OpenAI embedding model (1536 dims).
	ModelTextEmbeddingAda002 = "text-embedding-ada-002"
)

Well-known OpenAI embedding model names.

View Source 
const DefaultHFBaseURL = "https://huggingface.co"

DefaultHFBaseURL is the default base for HuggingFace model downloads. It can be overridden per-cache (CacheBaseURL) to point at a mirror or a local file server for offline use.

View Source 
const DefaultHFFile = "onnx/model.onnx"

DefaultHFFile is the default ONNX file within a HuggingFace repo.

View Source 
const DefaultHFRepo = "sentence-transformers/all-MiniLM-L6-v2"

DefaultHFRepo is the default HuggingFace repo used when none is given: the canonical sentence-transformers MiniLM ONNX export.

Variables

This section is empty.

Functions

func AutoDimension 

func AutoDimension(ctx context.Context, e Embedder, sample string) (int, error)

AutoDimension determines an embedder's output dimension by embedding a sample text. Useful for providers whose dimension is not well-known up front (e.g. arbitrary Ollama models).

func BundledTokenizerNames 

func BundledTokenizerNames() []string

BundledTokenizerNames returns the model names that ship with a bundled tokenizer, in sorted order.

func CosineSimilarity 

func CosineSimilarity(a, b []float32) float64

CosineSimilarity computes the cosine similarity between two vectors.

func HuggingFaceURL 

func HuggingFaceURL(repo, file string) string

HuggingFaceURL builds the download URL for an ONNX file in a HuggingFace repo using the default base URL.

func LoadHFModel 

func LoadHFModel(ctx context.Context, cache *ModelCache, repo, file string) (*onnx.Model, error)

LoadHFModel downloads (or reuses a cached copy of) the ONNX model for the given HuggingFace repo and returns a ready-to-run *onnx.Model. repo and file may be empty to use the defaults (sentence-transformers all-MiniLM-L6-v2, onnx/model.onnx). cache may be nil, in which case the model is fetched without caching.

Types

type CachingEmbedder 

type CachingEmbedder struct {
	// contains filtered or unexported fields
}

CachingEmbedder wraps another embedder with an embedding cache to avoid redundant (and often paid) API calls for repeated texts.

func NewCachingEmbedder 

func NewCachingEmbedder(inner Embedder, c *cache.EmbeddingCache, ttl time.Duration) *CachingEmbedder

NewCachingEmbedder wraps inner with the given cache. A non-positive TTL falls back to the cache package's default TTL.

func (*CachingEmbedder) Dimension 

func (c *CachingEmbedder) Dimension() int

Dimension returns the wrapped embedder's dimension.

func (*CachingEmbedder) Embed 

func (c *CachingEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed returns the cached vector when available, otherwise delegates to the wrapped embedder and stores the result.

func (*CachingEmbedder) EmbedBatch 

func (c *CachingEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch serves cached vectors where available and fetches only the missing texts from the wrapped embedder in a single batch call.

func (*CachingEmbedder) Stats 

func (c *CachingEmbedder) Stats() cache.CacheStats

Stats returns the underlying cache statistics merged with the hit/miss counts observed by this embedder (the LRU layer itself does not track lookups).

type CohereConfig 

type CohereConfig struct {
	// APIKey is the Cohere API key (required).
	APIKey string

	// Model is the embedding model name (required).
	Model string

	// BaseURL overrides the API endpoint (useful for proxies and tests).
	// Defaults to https://api.cohere.ai.
	BaseURL string

	// InputType is the embedding input type.
	// Defaults to InputTypeSearchDocument.
	InputType string

	// Truncation is the truncation strategy ("NONE", "START", "END").
	// Empty string omits the field (Cohere treats it as "NONE").
	Truncation string

	// Dimension optionally overrides the expected output dimension and is
	// used to validate responses. Zero means "use the model's native
	// dimension".
	Dimension int

	// BatchSize is the maximum number of texts per API request.
	// EmbedBatch splits larger inputs automatically.
	// Defaults to 96 (the Cohere API maximum).
	BatchSize int

	// Retry configures retry and backoff behavior.
	// Zero value uses DefaultRetryConfig.
	Retry RetryConfig

	// HTTPClient overrides the HTTP client. Defaults to a client with a
	// 30s timeout.
	HTTPClient *http.Client

	// Timeout is the request timeout when HTTPClient is nil.
	// Defaults to 30s.
	Timeout time.Duration
}

CohereConfig configures a CohereEmbedder.

type CohereEmbedder 

type CohereEmbedder struct {
	// contains filtered or unexported fields
}

CohereEmbedder embeds text using the Cohere embeddings API.

func NewCohereEmbedder 

func NewCohereEmbedder(cfg CohereConfig) (*CohereEmbedder, error)

NewCohereEmbedder creates a new CohereEmbedder, validating the configuration (API key, model, input type, and truncation strategy).

func (*CohereEmbedder) Dimension 

func (e *CohereEmbedder) Dimension() int

Dimension returns the embedding dimension of this embedder.

func (*CohereEmbedder) Embed 

func (e *CohereEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed converts a single text string into an embedding vector.

func (*CohereEmbedder) EmbedBatch 

func (e *CohereEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch converts multiple texts into embedding vectors, automatically splitting the input into API-sized batches.

type Embedder 

type Embedder interface {
	// Embed converts a single text string into a float32 embedding vector.
	Embed(ctx context.Context, text string) ([]float32, error)

	// EmbedBatch converts multiple text strings into embedding vectors.
	// Implementations may optimize batch processing for better performance.
	EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

	// Dimension returns the dimension of the embedding vectors produced by this embedder.
	Dimension() int
}

Embedder defines the interface for converting text into vector embeddings.

type MockEmbedder 

type MockEmbedder struct {
	// contains filtered or unexported fields
}

MockEmbedder is a simple embedder that produces deterministic pseudo-random vectors. Useful for testing and development without external dependencies.

func NewMockEmbedder 

func NewMockEmbedder(dim int) *MockEmbedder

NewMockEmbedder creates a MockEmbedder with the given dimension.

func (*MockEmbedder) Dimension 

func (m *MockEmbedder) Dimension() int

Dimension returns the embedding dimension.

func (*MockEmbedder) Embed 

func (m *MockEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed generates a deterministic embedding based on the text content.

func (*MockEmbedder) EmbedBatch 

func (m *MockEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch embeds multiple texts.

type MockMultiModal 

type MockMultiModal struct {
	// contains filtered or unexported fields
}

MockMultiModal is a deterministic multi-modal embedder for tests and offline use. Vectors are derived from FNV-1a hashes so that:

  • identical inputs produce identical vectors (cosine 1),
  • distinct inputs are near-orthogonal in expectation,
  • text and image hashes share the same space, so a query "about" content that reuses the same seed phrase matches deterministically (see SeedForText for building such pairs).

func NewMockMultiModal 

func NewMockMultiModal(dim int) *MockMultiModal

NewMockMultiModal creates a mock multi-modal embedder of the given dimension (min 8).

func (*MockMultiModal) Dimension 

func (m *MockMultiModal) Dimension() int

Dimension returns the configured dimension.

func (*MockMultiModal) EmbedImage 

func (m *MockMultiModal) EmbedImage(_ context.Context, data []byte, mimeType string) ([]float32, error)

EmbedImage embeds image bytes deterministically.

func (*MockMultiModal) EmbedText 

func (m *MockMultiModal) EmbedText(_ context.Context, text string) ([]float32, error)

EmbedText embeds text deterministically.

func (*MockMultiModal) SeedForText 

func (m *MockMultiModal) SeedForText(text string) string

SeedForText returns the canonical seed string for a text embedding. Tests use it to assert exact vector equality.

type Modality 

type Modality string

Modality labels the kind of content stored in a multi-modal index. 

const (
	// ModalityText is plain text content.
	ModalityText Modality = "text"
	// ModalityImage is raw image content.
	ModalityImage Modality = "image"
)

type ModelCache 

type ModelCache struct {
	// contains filtered or unexported fields
}

ModelCache is a small on-disk cache for ONNX model files. Entries are keyed by the SHA-256 of the source URL and validated by content hash (the filename), so identical content from different URLs shares a file only when explicitly requested; the primary key is the URL hash.

The cache is safe for concurrent use. Downloads are deduplicated: while one goroutine is fetching a URL, others wait for the same result.

func NewModelCache 

func NewModelCache(dir string, ttl time.Duration) (*ModelCache, error)

NewModelCache creates a cache rooted at dir. ttl is the maximum age of a cached file before it is re-downloaded (0 means cache forever). The file is stored under the name sha256(url).onnx.

func (*ModelCache) Get 

func (c *ModelCache) Get(ctx context.Context, rawURL string) (string, error)

Get returns the on-disk path for the model at rawURL, downloading it if the cache is empty or the entry is older than the TTL. The returned path is suitable for onnx.LoadFile.

func (*ModelCache) HFURL 

func (c *ModelCache) HFURL(repo, file string) string

HFURL builds the download URL for an ONNX file in a HuggingFace repo using this cache's base URL: <base>/<repo>/resolve/main/<file>.

func (*ModelCache) Path 

func (c *ModelCache) Path(rawURL string) string

Path returns the on-disk path a given URL would be cached to.

func (*ModelCache) SetBaseURL 

func (c *ModelCache) SetBaseURL(base string)

SetBaseURL overrides the default HuggingFace base URL (for example to point at a mirror or a local file server for offline use). An empty string restores the default.

type MultiModalEmbedder 

type MultiModalEmbedder interface {
	// EmbedText converts text into a vector in the shared space.
	EmbedText(ctx context.Context, text string) ([]float32, error)

	// EmbedImage converts raw image bytes (with MIME type) into a
	// vector in the shared space.
	EmbedImage(ctx context.Context, data []byte, mimeType string) ([]float32, error)

	// Dimension returns the shared embedding dimension.
	Dimension() int
}

MultiModalEmbedder embeds heterogeneous content (text and images) into a single shared vector space, enabling cross-modal retrieval: text queries can match stored images and vice versa. Real providers include CLIP-family models; the interface keeps them pluggable.

type OllamaConfig 

type OllamaConfig struct {
	// Model is the Ollama embedding model name (required), e.g.
	// "all-minilm-l6-v2" or any model pulled via `ollama pull`.
	Model string

	// BaseURL is the Ollama server URL.
	// Defaults to http://localhost:11434.
	BaseURL string

	// Dimension optionally fixes the expected output dimension (used to
	// validate responses). For models not in ollamaKnownDimensions this is
	// required — or call DetectDimension before using the embedder with a
	// store, since stores pin the dimension at construction time.
	Dimension int

	// BatchSize is the maximum number of texts per API request.
	// EmbedBatch splits larger inputs automatically. Defaults to 32.
	BatchSize int

	// Retry configures retry and backoff behavior.
	// Zero value uses DefaultRetryConfig.
	Retry RetryConfig

	// HTTPClient overrides the HTTP client. Defaults to a client with a
	// 60s timeout.
	HTTPClient *http.Client

	// Timeout is the request timeout when HTTPClient is nil.
	// Defaults to 60s.
	Timeout time.Duration
}

OllamaConfig configures an OllamaEmbedder.

type OllamaEmbedder 

type OllamaEmbedder struct {
	// contains filtered or unexported fields
}

OllamaEmbedder embeds text using a local model served by an Ollama instance. This is the zero-CGO path for local embedding models: Ollama handles model download, caching, and CPU inference.

func NewOllamaEmbedder 

func NewOllamaEmbedder(cfg OllamaConfig) (*OllamaEmbedder, error)

NewOllamaEmbedder creates a new OllamaEmbedder.

func (*OllamaEmbedder) DetectDimension 

func (e *OllamaEmbedder) DetectDimension(ctx context.Context, sample string) (int, error)

DetectDimension embeds a sample text and returns the resulting vector dimension, updating the embedder's known dimension in the process. Use it before constructing a store with a model whose dimension is not well-known.

func (*OllamaEmbedder) Dimension 

func (e *OllamaEmbedder) Dimension() int

Dimension returns the embedding dimension of this embedder. It may return 0 for models of unknown dimension until DetectDimension or the first Embed call has resolved it.

func (*OllamaEmbedder) Embed 

func (e *OllamaEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed converts a single text string into an embedding vector.

func (*OllamaEmbedder) EmbedBatch 

func (e *OllamaEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch converts multiple texts into embedding vectors, automatically splitting the input into batches.

type OnnxEmbedder 

type OnnxEmbedder struct {
	// contains filtered or unexported fields
}

OnnxEmbedder is a zero-network embedder backed by a pure-Go ONNX inference runtime. It runs sentence-transformer (or similar) ONNX exports locally.

func NewOnnxEmbedder 

func NewOnnxEmbedder(cfg OnnxEmbedderConfig) (*OnnxEmbedder, error)

NewOnnxEmbedder creates an embedder from a loaded ONNX model.

func NewOnnxEmbedderFile 

func NewOnnxEmbedderFile(path string, cfg OnnxEmbedderConfig) (*OnnxEmbedder, error)

NewOnnxEmbedderFile loads an ONNX model from disk and creates an embedder.

func (*OnnxEmbedder) Dimension 

func (e *OnnxEmbedder) Dimension() int

Dimension returns the output dimension, probing the model once with a minimal text if it was not pinned in the config.

func (*OnnxEmbedder) Embed 

func (e *OnnxEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed converts a single text into an embedding vector.

func (*OnnxEmbedder) EmbedBatch 

func (e *OnnxEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch converts multiple texts into embedding vectors, one per text. Sequences are tokenized sequentially (cheap) and then executed in parallel by the ONNX runtime, with the worker count controlled by BatchConcurrency (0 = auto).

type OnnxEmbedderConfig 

type OnnxEmbedderConfig struct {
	// Model is a loaded ONNX model (required).
	Model *onnx.Model

	// Tokenize converts each text into the model's named inputs (required).
	Tokenize TokenizerFunc

	// Output names the model output to read as the embedding. When empty,
	// the model's last declared output is used.
	Output string

	// Normalize L2-normalizes the resulting vectors. Most
	// sentence-transformer models expect normalized inputs for cosine
	// similarity search.
	Normalize bool

	// Dimension optionally pins the output dimension, skipping the lazy
	// probe run that Dimension() would otherwise perform.
	Dimension int

	// BatchConcurrency caps how many sequences EmbedBatch executes in
	// parallel. Zero (the default) auto-selects a worker count from the
	// available CPUs (capped at 8). Values <= 0 are treated as the
	// default; a value of 1 forces sequential execution. Note that peak
	// memory scales linearly with concurrency, since each in-flight
	// sequence holds its full intermediate tensor state.
	BatchConcurrency int
}

OnnxEmbedderConfig configures an OnnxEmbedder.

type OpenAIConfig 

type OpenAIConfig struct {
	// APIKey is the OpenAI API key (required).
	APIKey string

	// Model is the embedding model name (required).
	Model string

	// BaseURL overrides the API endpoint (useful for proxies and tests).
	// Defaults to https://api.openai.com/v1.
	BaseURL string

	// Dimension optionally requests a reduced output dimension (the
	// Matryoshka "dimension" parameter) and is used to validate responses.
	// It must not exceed the model's native dimension. Zero means "use the
	// model's native dimension".
	Dimension int

	// BatchSize is the maximum number of texts per API request.
	// EmbedBatch splits larger inputs automatically.
	// Defaults to 100 and is capped at 2048.
	BatchSize int

	// Retry configures retry and backoff behavior.
	// Zero value uses DefaultRetryConfig.
	Retry RetryConfig

	// HTTPClient overrides the HTTP client. Defaults to a client with a
	// 30s timeout.
	HTTPClient *http.Client

	// Timeout is the request timeout when HTTPClient is nil.
	// Defaults to 30s.
	Timeout time.Duration
}

OpenAIConfig configures an OpenAIEmbedder.

type OpenAIEmbedder 

type OpenAIEmbedder struct {
	// contains filtered or unexported fields
}

OpenAIEmbedder embeds text using the OpenAI embeddings API.

func NewOpenAIEmbedder 

func NewOpenAIEmbedder(cfg OpenAIConfig) (*OpenAIEmbedder, error)

NewOpenAIEmbedder creates a new OpenAIEmbedder, validating the configuration (API key, model, and dimension constraints).

func (*OpenAIEmbedder) Dimension 

func (e *OpenAIEmbedder) Dimension() int

Dimension returns the embedding dimension of this embedder.

func (*OpenAIEmbedder) Embed 

func (e *OpenAIEmbedder) Embed(ctx context.Context, text string) ([]float32, error)

Embed converts a single text string into an embedding vector.

func (*OpenAIEmbedder) EmbedBatch 

func (e *OpenAIEmbedder) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch converts multiple texts into embedding vectors, automatically splitting the input into API-sized batches.

type Pipeline 

type Pipeline struct {
	// contains filtered or unexported fields
}

Pipeline chains embedders and returns the first successful result, enabling failover between providers (e.g. a paid API with a local model as backup). All embedders must produce the same dimension, otherwise a store's index would silently mix incompatible vectors when failover kicks in.

func NewPipeline 

func NewPipeline(embedders ...Embedder) (*Pipeline, error)

NewPipeline creates a Pipeline from the given embedders, tried in order.

func (*Pipeline) Dimension 

func (p *Pipeline) Dimension() int

Dimension returns the embedding dimension of this pipeline (the first embedder's dimension).

func (*Pipeline) Embed 

func (p *Pipeline) Embed(ctx context.Context, text string) ([]float32, error)

Embed converts a single text using the first embedder that succeeds.

func (*Pipeline) EmbedBatch 

func (p *Pipeline) EmbedBatch(ctx context.Context, texts []string) ([][]float32, error)

EmbedBatch converts multiple texts using the first embedder that succeeds for the whole batch.

type RetryConfig 

type RetryConfig struct {
	// MaxAttempts is the total number of attempts, including the first.
	// Zero uses the default (3).
	MaxAttempts int

	// InitialBackoff is the base delay before the first retry.
	// Zero uses the default (500ms).
	InitialBackoff time.Duration

	// MaxBackoff caps the exponential backoff delay.
	// Zero uses the default (8s).
	MaxBackoff time.Duration
}

RetryConfig controls retry behavior for embedder HTTP requests.

func DefaultRetryConfig 

func DefaultRetryConfig() RetryConfig

DefaultRetryConfig returns sensible retry defaults.

type TokenizerFunc 

type TokenizerFunc func(text string) (map[string]*onnx.Tensor, error)

TokenizerFunc converts a text into the named input tensors an ONNX embedding model expects (e.g. "input_ids", "attention_mask"). Tokenization is deliberately dependency-injected: the ONNX runtime executes the model, while the caller supplies the model-specific tokenizer.

func BundledTokenizer 

func BundledTokenizer(modelName string, m *onnx.Model) (TokenizerFunc, error)

BundledTokenizer returns a ready-to-use TokenizerFunc for a known ONNX model name. The function tokenizes text with the model's exact configuration (lowercasing, max length, special tokens) and emits only the input tensors the model declares, padded to the model's maximum sequence length. The vocab is embedded, so this works fully offline.

type Wordpiece 

type Wordpiece struct {
	// contains filtered or unexported fields
}

Wordpiece is a BERT-style WordPiece tokenizer (compatible with HuggingFace's BertTokenizer for the supported options). It is safe for concurrent use: all state is read-only after construction.

func NewWordpiece 

func NewWordpiece(cfg WordpieceConfig) (*Wordpiece, error)

NewWordpiece builds a tokenizer from the given config. It returns an error when required tokens ([CLS], [SEP], [PAD]) are missing from the vocab.

func (*Wordpiece) AsTokenizerFunc 

func (w *Wordpiece) AsTokenizerFunc(m *onnx.Model) TokenizerFunc

AsTokenizerFunc returns a TokenizerFunc bound to this tokenizer and model.

func (*Wordpiece) Encode 

func (w *Wordpiece) Encode(text string) (ids, mask, types []int)

Encode converts text into token ids, an attention mask and token type ids, each of the same length (padded to cfg.PadTo when set). The CLS and SEP tokens (when configured) are always preserved on truncation.

func (*Wordpiece) FeedsForModel 

func (w *Wordpiece) FeedsForModel(m *onnx.Model, text string) (map[string]*onnx.Tensor, error)

FeedsForModel converts the tokenizer's Encode output into the named input tensors the model actually declares (input_ids, attention_mask, token_type_ids). This is the standard bridge from a Wordpiece tokenizer to an ONNX model's feed inputs.

func (*Wordpiece) TokenID 

func (w *Wordpiece) TokenID(tok string) int

TokenID returns the id of a vocab token, or -1 when unknown.

func (*Wordpiece) VocabSize 

func (w *Wordpiece) VocabSize() int

VocabSize returns the number of tokens in the vocabulary.

type WordpieceConfig 

type WordpieceConfig struct {
	// Vocab is the vocab file content, one token per line, in id order.
	// When empty, the embedded BERT-uncased vocabulary is used.
	Vocab string

	// DoLowerCase lowercases the input before tokenization (BERT-uncased
	// models set this to true).
	DoLowerCase bool

	// MaxLength caps the number of tokens produced by Encode, preserving
	// the leading CLS and trailing SEP tokens when present.
	MaxLength int

	// PadTo pads every Encode output to this length with the PadTokenID
	// (and a 0 in the attention mask). When <= 0 the output is unpadded.
	// Most ONNX exports expect fixed-shape inputs, so callers should set
	// this to the model's maximum sequence length.
	PadTo int

	// CLSToken and SEPToken are prepended/appended around the tokens.
	// Empty disables them (e.g. for models that expect raw input).
	CLSToken string
	SEPToken string

	// PadTokenID is used for padding. When <= 0, the id of "[PAD]" in the
	// vocabulary is used.
	PadTokenID int

	// TokenTypeIDs, when non-nil, is used verbatim for the token_type_ids
	// tensor (zero-filled beyond its length). When nil, all zeros are
	// produced.
	TokenTypeIDs []int
}

WordpieceConfig configures a BERT-style WordPiece tokenizer.

Source Files

View all Source files

Directories

Path Synopsis
Package onnx is a minimal, pure-Go ONNX model loader and interpreter for inference-only use.
 Package onnx is a minimal, pure-Go ONNX model loader and interpreter for inference-only use.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.