Implementation:Ollama Ollama Imagegen Safetensors Extractor

Knowledge Sources	Ollama
Domains	Image Generation, Model Loading
Last Updated	2025-02-15 00:00 GMT

Overview

Extracts individual tensors from safetensors files with io.Reader interfaces for streaming writes, enabling per-tensor blob creation without loading entire files into memory.

Description

The extractor.go file provides TensorExtractor and TensorData for streaming tensor extraction from safetensors files. TensorExtractor parses the JSON header to build a tensor index, then provides io.SectionReader-based access to each tensor's raw bytes. TensorData holds metadata (name, dtype, shape, size) and provides SafetensorsReader which wraps the raw data in a minimal safetensors format (8-byte header size + JSON header + data) for native mlx_load_safetensors compatibility. BuildPackedSafetensorsReader creates a streaming reader for multiple tensors packed into a single safetensors blob, used for expert group packing. ExtractRawFromSafetensors strips the header from a safetensors reader to get raw tensor bytes.

Usage

Used during model import (ollama create) to extract individual tensors from HuggingFace safetensors files and create per-tensor blobs for Ollama's storage format.

Code Reference

Source Location

• Repository: Ollama
• File: x/imagegen/safetensors/extractor.go
• Lines: 1-244

Signature

type TensorExtractor struct {
	file       *os.File
	dataOffset int64
	header     map[string]tensorInfo
}

type TensorData struct {
	Name   string
	Dtype  string
	Shape  []int32
	Size   int64
	reader *io.SectionReader
}

func OpenForExtraction(path string) (*TensorExtractor, error)
func (te *TensorExtractor) ListTensors() []string
func (te *TensorExtractor) GetTensor(name string) (*TensorData, error)
func (td *TensorData) Reader() io.Reader
func (td *TensorData) SafetensorsReader() io.Reader
func (td *TensorData) SafetensorsSize() int64
func BuildPackedSafetensorsReader(tensors []*TensorData) io.Reader
func ExtractRawFromSafetensors(r io.Reader) ([]byte, error)

Import

import "github.com/ollama/ollama/x/imagegen/safetensors"

I/O Contract

Inputs

Name	Type	Required	Description
path	string	Yes	Path to .safetensors file
name	string	Yes	Tensor name to extract

Outputs

Name	Type	Description
*TensorData	*TensorData	Tensor metadata with streaming reader
io.Reader	io.Reader	Streaming reader for safetensors-wrapped tensor

Usage Examples

extractor, err := safetensors.OpenForExtraction("model.safetensors")
if err != nil {
    return err
}
defer extractor.Close()

names := extractor.ListTensors()
for _, name := range names {
    td, _ := extractor.GetTensor(name)
    // Stream as safetensors blob
    reader := td.SafetensorsReader()
    // Write to blob storage...
}

Related Pages

• Principle:Ollama_Ollama_ImageGeneration