transformer

package
v0.3.5 Latest Latest
Warning

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

 Go to latest Published: Apr 16, 2026 License: Apache-2.0 Imports: 20 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"`
	IntermediateSize  int      `json:"intermediate_size"`
	NumKeyValueHeads  int      `json:"num_key_value_heads"`
	RMSNormEps        float64  `json:"rms_norm_eps"`
	LayerNormEps      float64  `json:"layer_norm_eps"`

	// RoPE Positional Embedder
	RoPETheta   float64     `json:"rope_theta"`
	RoPEScaling RoPEScaling `json:"rope_scaling"`

	// 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"`
	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
	// 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(ctx *context.Context, tokens, mask *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.
  • mask: shaped [batchSize, seqLen] with the mask, where 1 means valid token and 0 means padding. It can be nil, in which case not mask is set.

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, mask *graph.Node) *graph.Node

ApplySentencePooling applies the configured pooling function to the hidden states.

- hiddenStates: [batchSize, seqLen, hiddenSize] - mask: nil or [batchSize, seqLen] of dtype Bool.

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(ctx *context.Context) *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) ForwardGraph added in v0.3.5 

func (m *Model) ForwardGraph(ctx *context.Context, tokens, mask *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) LoadContext 

func (m *Model) LoadContext(backend backends.Backend, ctx *context.Context) error

LoadContext uses models/safetensors to load the variables of the model into a context.

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) 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(ctx *context.Context, tokens, mask *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.
  • 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.

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 backends.Backend, ctx *context.Context) (*context.Exec, error)

SingleSentenceEmbeddingExec returns a context.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 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.