Documentation
¶
Overview ¶
Package kaggle loads Gemma weights into tensors along with the matching metadata, after they have been downloaded from kaggle and converted using the included cmd/convert_checkpoint.py.
Index ¶
- Constants
- func ParamNames(metadata *trees.Tree[*Metadata]) *trees.Tree[string]
- func PyReadAggregate(checkpointDir string) (results any, err error)
- func PyReadParamInfo(checkpointDir string) *trees.Tree[*PyParamInfo]
- func ReadConvertedWeights(ctx *context.Context, checkpointDir string) error
- func ReadConvertedWeightsToTree(checkpointDir string) (tree *trees.Tree[*tensors.Tensor], err error)
- func ReadMetadata(checkpointDir string) (tree *trees.Tree[*Metadata], err error)
- func UploadWeightsToContext(ctx *context.Context, weights *trees.Tree[*tensors.Tensor])
- type Metadata
- type MetadataKeyType
- type PyParamInfo
Constants ¶
const ( MetadataFileName = "_METADATA" KeyUseZarr3 = "use_zarr3" KeyTreeMetadata = "tree_metadata" KeyKeyMetadata = "key_metadata" KeyValueMetadata = "value_metadata" )
const ( AggregateFileName = "checkpoint" // OCDBTManifestFileName indicates usage of "Orbax Consistent Distributed Backend Tree" (OCDBT). OCDBTManifestFileName = "manifest.ocdbt" )
Variables ¶
This section is empty.
Functions ¶
func ParamNames ¶
ParamNames convert metadata to the paramNames (?? not sure where in Gemma this is sued)
func PyReadAggregate ¶
PyReadAggregate of Python checkpoint. Not used by Gemma v2.
func PyReadParamInfo ¶
func PyReadParamInfo(checkpointDir string) *trees.Tree[*PyParamInfo]
func ReadConvertedWeights ¶
ReadConvertedWeights from checkpointDir (under the "raw/" subdirectory). It will read the weights and shape converted by the `convert_checkpoint.py` script (see github.com/gomlx/gemma repository, under cmd/convert_checkpoint.py) and set them in the given context, under its current scope.
func ReadConvertedWeightsToTree ¶
func ReadConvertedWeightsToTree(checkpointDir string) (tree *trees.Tree[*tensors.Tensor], err error)
ReadConvertedWeightsToTree from checkpointDir (under the "raw/" subdirectory). It will read the weights and shape converted by the `convert_checkpoint.py` script (see github.com/gomlx/gemma repository, under cmd/convert_checkpoint.py)
It returns a tree of tensors, with the path matching those of the original Jax checkpoint.
func ReadMetadata ¶
ReadMetadata returns the metadata loaded from the given directory in the form of a tree.
func UploadWeightsToContext ¶
UploadWeightsToContext creates variables corresponding to the weights. It returns the ctx given, with the variables set.
It's tightly coupled with the model building functions in this package. Meaning the modeling must match the naming here.
Types ¶
type Metadata ¶
type Metadata struct {
KeyPath []string
KeyToKeyType map[string]MetadataKeyType
ValueType string
SkipDeserialize bool
}
Metadata of one checkpoint entry (usually weights or embedding tables of model)
type MetadataKeyType ¶
type MetadataKeyType int
const ( KeyTypeSequence MetadataKeyType = 1 KeyTypeDict = 2 )
Source Files