transformer

package
v0.4.0 Latest Latest
Warning

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

 Go to latest Published: Jul 21, 2026 License: Apache-2.0 Imports: 23 Imported by: 0

Documentation

Index

Constants

View Source 
const DefaultQueryPrompt = "Instruct: Given a query, retrieve documents that answer the query \nQuery: "

Variables

This section is empty.

Functions

This section is empty.

Types

type Config 

type Config struct {
	Architectures         []string `json:"architectures"`
	HiddenSize            int      `json:"hidden_size"`
	NumHiddenLayers       int      `json:"num_hidden_layers"`
	NumAttentionHeads     int      `json:"num_attention_heads"`
	HeadDim               int      `json:"head_dim"`
	GlobalHeadDim         int      `json:"global_head_dim"`
	IntermediateSize      int      `json:"intermediate_size"`
	NumKeyValueHeads      int      `json:"num_key_value_heads"`
	NumKVSharedLayers     int      `json:"num_kv_shared_layers"`
	RMSNormEps            float64  `json:"rms_norm_eps"`
	LayerNormEps          float64  `json:"layer_norm_eps"`
	AttentionLogitCap     float64  `json:"attention_logit_cap"`
	FinalLogitSoftcapping float64  `json:"final_logit_softcapping"`

	HiddenSizePerLayerInput int `json:"hidden_size_per_layer_input"`
	VocabSizePerLayerInput  int `json:"vocab_size_per_layer_input"`

	// RoPE Positional Embedder
	RoPETheta      float64               `json:"rope_theta"`
	RoPEScaling    RoPEScaling           `json:"rope_scaling"`
	RoPEParameters map[string]RoPEParams `json:"rope_parameters"`

	// RoPELocalBaseFreq is the Theta used by sliding attention layers (or so the AI says).
	RoPELocalBaseFreq float64 `json:"rope_local_base_freq"`

	HiddenActivation      string `json:"hidden_activation"`
	HiddenAct             string `json:"hidden_act"`
	MaxPositionEmbeddings int    `json:"max_position_embeddings"`
	ModelType             string `json:"model_type"`
	TorchDtype            string `json:"torch_dtype"`
	DType                 string `json:"dtype"`
	VocabSize             int    `json:"vocab_size"`
	PadTokenID            int    `json:"pad_token_id"`

	// SlidingWindow is the size of the sliding window for sliding attention layers.
	SlidingWindow int `json:"sliding_window"`

	// LayerTypes: known values are "full_attention", "sliding_attention"
	LayerTypes []string       `json:"layer_types"`
	Extra      map[string]any `json:"-"`
}

Config represents standard HuggingFace config.json.

func (*Config) UnmarshalJSON 

func (c *Config) UnmarshalJSON(data []byte) error

type Model 

type Model struct {
	Repo *hub.Repo

	Config                    Config
	SentenceTransformerConfig *SentenceTransformerConfig
	Modules                   []ModuleConfig
	PoolingConfig             *PoolingConfig

	KVCache *mltransformer.KVCache
	// contains filtered or unexported fields
}

Model holds all the configuration loaded about the model.

func LoadModel 

func LoadModel(repo *hub.Repo) (*Model, error)

LoadModel loads the configurations into the Model struct. It loads config.json as mandatory, and gracefully handles the optional presence of sentence_transformer and pooling configurations.

func (*Model) AllLayers added in v0.3.5 

func (m *Model) AllLayers(scope *model.Scope, tokens, seqLen *graph.Node) (lastLayer *graph.Node, allLayers []*graph.Node)

AllLayers takes the input tokens and creates the GoMLX forward graph for the transformer model, returning the last layer and all the intermediate layers.

Inputs:

  • tokens: shaped [batchSize, seqLen] with the tokens, including padding.
  • seqLen: shaped [batchSize] or scalar, with the number of non-padding tokens. It can be nil.

It returns:

  • lastLayer: the final hidden state of the last layer, shaped [batchSize, seqLen,hiddenSize].
  • allLayers: the input to the first layer and the output of each layer. It follows the HuggingFace convention, where the allLayers[0] is the input to the first attention layer, and the following nodes in allLayers are the outputs of all NumHiddenLayers attention layers.

func (*Model) ApplySentencePooling 

func (m *Model) ApplySentencePooling(hiddenStates, seqLen *graph.Node) *graph.Node

ApplySentencePooling applies the configured pooling function to the hidden states.

- hiddenStates: [batchSize, seqLen, hiddenSize] - seqLen: nil or [batchSize] of type Int32.

Returns [batchSize, hiddenSize]

func (*Model) BuildPrompt added in v0.3.5 

func (m *Model) BuildPrompt(sentence, promptName string) string

BuildPrompt builds the full sentence prompt, based on a promptName, an index to the list of prompts in the SentenceTransformerConfig. If the code is not found, it attempts to use the default one. If a default one is not defined, it returns the original sentence.

func (*Model) BuildQueryPrompt added in v0.3.5 

func (m *Model) BuildQueryPrompt(query, promptName string) string

BuildQueryPrompt builds the full query prompt, based on a promptName. It's exactly like BuildPrompt but if a default prompt doesn't exist it uses "Instruct: Given a query, retrieve documents that answer the query \nQuery: " as a prompt prefix.

func (*Model) CreateGoMLXModel added in v0.3.5 

func (m *Model) CreateGoMLXModel(scope *model.Scope) *mltransformer.Model

CreateGoMLXModel initializes the base GoMLX ml_transformer.Model configuration using the loaded fields.

Usually, you won't call this directly, instead you would use the Model.SentenceEmbedding or Model.AllLayers. But you can use it for something custom.

It takes the context ctx with the loaded variables.

func (*Model) Description 

func (m *Model) Description() string

Description returns a string summarizing the model architecture, size, and loaded configurations.

func (*Model) Forward added in v0.4.0 

func (m *Model) Forward(
	scope *model.Scope,
	tokenIds *graph.Node,
	positionIds *graph.Node,
	seqLen *graph.Node,
	attentionMask *graph.Node,
	cache mltransformer.KVCacheNodes,
) (logits *graph.Node, updatedCache mltransformer.KVCacheNodes)

Forward performs the forward pass of the model. It returns the logits of the vocabulary projection (typically shaped [batchSize, seqLen, vocabSize]) and the updated KV Cache.

func (*Model) ForwardGraph added in v0.3.5 

func (m *Model) ForwardGraph(scope *model.Scope, tokens, seqLen *graph.Node) *graph.Node

ForwardGraph takes the input tokens and creates the GoMLX graph for the model. It returns the final sentence embeddings if appropriate, otherwise just final hidden states.

  • tokens: shaped [batchSize, seqLen] with the tokens, including padding.
  • mask: shaped [batchSize, seqLen] with the mask, where 1 means valid token and 0 means padding. It can be nil, in which case no mask is used and it's assumed all elements of the sentence are used.

If the model was trained as an embedding model (e.g. sentence-transformers), it will return the sentence embeddings, usually as [batchSize, embedSize]. Otherwise, it will return the final hidden states of all layers, usually as [batchSize, seqLen, hiddenSize].

func (*Model) GetTaskPrompt added in v0.3.5 

func (m *Model) GetTaskPrompt(taskCode string) string

GetTaskPrompt returns the prompt string for the given task code. Returns an empty string if the task code is not found or no prompts are registered.

func (*Model) GetTokenizer added in v0.3.5 

func (m *Model) GetTokenizer() (tokenizers.Tokenizer, error)

GetTokenizer returns the tokenizer for the model if it exists, or if it hasn't been set yet, attempts to create the default tokenizer for the repo.

func (*Model) KVCacheConfig added in v0.4.0 

func (m *Model) KVCacheConfig() *mltransformer.KVCache

KVCacheConfig returns the KVCache configuration struct for the model. The user can modify this configuration as needed.

func (*Model) LoadStore added in v0.4.0 

func (m *Model) LoadStore(backend compute.Backend, store *model.Store) error

LoadStore uses models/safetensors to load the variables of the model into a GoMLX's model.Store.

If a backend is provided (not nil), the variables are immediately loaded into the backend device #0, saving host memory space or accelerating the loading in some cases.

For distributed execution, better to leave backend and nil, and let the executor decide on which devices to place the variables.

func (*Model) LoadStoreFiltered added in v0.4.0 

func (m *Model) LoadStoreFiltered(backend compute.Backend, store *model.Store, filterFn func(path string) bool) error

LoadStoreFiltered is like LoadStore but takes a filter function. Only variables where filterFn(fullPath) returns true will be loaded. Other tensors will be finalized immediately to conserve memory.

func (*Model) RegisteredPromptTasks added in v0.3.5 

func (m *Model) RegisteredPromptTasks() []string

RegisteredPromptTasks returns a list of all task codes for which prompts are registered.

func (*Model) SentenceEmbeddingGraph 

func (m *Model) SentenceEmbeddingGraph(scope *model.Scope, tokens, seqLen *graph.Node) *graph.Node

SentenceEmbeddingGraph builds the equivalent of the sentence_transformers pipeline. It uses AllLayers for the base model, and then applies pooling and normalization layers sequentially according to the modules.json configuration.

  • tokens: shaped [batchSize, seqLen] with the tokens, including padding.
  • seqLen: shaped [batchSize] or scalar, with the number of non-padding tokens. It can be nil.

It returns the final pooled embedding (usually [batchSize, embedSize]) for the sentence.

func (*Model) SetTokenizer added in v0.3.5 

func (m *Model) SetTokenizer(tok tokenizers.Tokenizer)

SetTokenizer sets the tokenizer of the Model.

func (*Model) Similarity added in v0.3.5 

func (m *Model) Similarity(queries, documents *Node) *Node

Similiarity across all queries and documents, based on the model similarity configuration.

- queries, documents: either [batchSize, embedDim] or for single query/documents [embedDim] - returns: [batchSizeQueries, batchSizeDocuments] matrix of similarities

func (*Model) SingleSentenceEmbeddingExec added in v0.3.5 

func (m *Model) SingleSentenceEmbeddingExec(backend compute.Backend, scope *model.Scope) (*model.Exec, error)

SingleSentenceEmbeddingExec returns a model.Exec that can be used to compute sentence embeddings. No padding, not bucketing, the exec takes as input a single sentence [seqLen] and returns the embedding [embedDim].

func (*Model) WithCausalMask 

func (m *Model) WithCausalMask(useCausalMask bool) *Model

WithCausalMask sets whether to use a causal mask in the attention layers.

Some models are trained with a causal mask, others without, but it is not documented in the usual model configuration. 

The default is to use a causal mask.

type ModuleConfig 

type ModuleConfig struct {
	Idx   int            `json:"idx"`
	Name  string         `json:"name"`
	Path  string         `json:"path"`
	Type  string         `json:"type"`
	Extra map[string]any `json:"-"`
}

ModuleConfig represents an entry in modules.json

func (*ModuleConfig) UnmarshalJSON 

func (m *ModuleConfig) UnmarshalJSON(data []byte) error

type PoolingConfig 

type PoolingConfig struct {
	PoolingModeMeanTokens bool           `json:"pooling_mode_mean_tokens"`
	PoolingModeClsToken   bool           `json:"pooling_mode_cls_token"`
	PoolingModeMaxTokens  bool           `json:"pooling_mode_max_tokens"`
	PoolingModeLastToken  bool           `json:"pooling_mode_lasttoken"`
	Extra                 map[string]any `json:"-"`
}

PoolingConfig represents 1_Pooling/config.json

func (*PoolingConfig) UnmarshalJSON 

func (c *PoolingConfig) UnmarshalJSON(data []byte) error

type RoPEParams added in v0.4.0 

type RoPEParams struct {
	RopeTheta           float64 `json:"rope_theta"`
	RopeType            string  `json:"rope_type"`
	PartialRotaryFactor float64 `json:"partial_rotary_factor"`
}

type RoPEScaling added in v0.3.5 

type RoPEScaling struct {
	Factor float64 `json:"factor"`

	// Type: "linear"
	Type string `json:"rope_type"`
}

type SentenceTransformerConfig 

type SentenceTransformerConfig struct {
	Extra             map[string]any    `json:"-"`
	Prompts           map[string]string `json:"prompts"`
	DefaultPromptName string            `json:"default_prompt_name"`
	SimilarityFnName  string            `json:"similarity_fn_name"`
}

SentenceTransformerConfig represents config_sentence_transformers.json (optional)

func (*SentenceTransformerConfig) UnmarshalJSON 

func (c *SentenceTransformerConfig) UnmarshalJSON(data []byte) error

Source Files

View all Source files

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.