ollama source for Momentry Core verification

x/safetensors/extractor.go

@@ -0,0 +1,283 @@
+package safetensors
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+	"sort"
+)
+
+// tensorInfo holds tensor metadata from safetensors headers.
+type tensorInfo struct {
+	Dtype       string  `json:"dtype"`
+	Shape       []int32 `json:"shape"`
+	DataOffsets [2]int  `json:"data_offsets"`
+}
+
+// TensorExtractor extracts individual tensors from a safetensors file.
+// It provides io.Reader interfaces for each tensor's raw data, enabling
+// streaming writes to blobs without loading entire tensors into memory.
+type TensorExtractor struct {
+	file       *os.File
+	dataOffset int64 // Start of tensor data region
+	header     map[string]tensorInfo
+}
+
+// TensorData holds tensor metadata and a reader for its raw bytes.
+type TensorData struct {
+	Name   string
+	Dtype  string
+	Shape  []int32
+	Size   int64
+	reader *io.SectionReader
+}
+
+// WithName returns a shallow copy of TensorData with a different logical tensor
+// name but the same underlying raw data reader.
+func (td *TensorData) WithName(name string) *TensorData {
+	if td == nil {
+		return nil
+	}
+	shape := make([]int32, len(td.Shape))
+	copy(shape, td.Shape)
+	return &TensorData{
+		Name:   name,
+		Dtype:  td.Dtype,
+		Shape:  shape,
+		Size:   td.Size,
+		reader: td.reader,
+	}
+}
+
+// Reader returns an io.Reader for the tensor's raw bytes.
+func (td *TensorData) Reader() io.Reader {
+	return td.reader
+}
+
+// safetensorsHeader builds the JSON header for a minimal safetensors blob
+// containing a single tensor keyed by its name.
+func (td *TensorData) safetensorsHeader() []byte {
+	header := map[string]any{
+		td.Name: tensorInfo{
+			Dtype:       td.Dtype,
+			Shape:       td.Shape,
+			DataOffsets: [2]int{0, int(td.Size)},
+		},
+	}
+	headerJSON, _ := json.Marshal(header)
+
+	// Pad header to 8-byte alignment
+	padding := (8 - len(headerJSON)%8) % 8
+	headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...)
+	return headerJSON
+}
+
+// SafetensorsReader returns a reader that outputs the tensor wrapped in
+// minimal safetensors format. This allows using mlx_load_safetensors on
+// individual tensor blobs for native zero-copy loading.
+// The tensor is keyed by its name in the safetensors header.
+func (td *TensorData) SafetensorsReader() io.Reader {
+	headerJSON := td.safetensorsHeader()
+
+	// Build header with size prefix
+	headerBuf := new(bytes.Buffer)
+	binary.Write(headerBuf, binary.LittleEndian, uint64(len(headerJSON)))
+	headerBuf.Write(headerJSON)
+
+	// Return multi-reader: header + tensor data
+	td.reader.Seek(0, io.SeekStart)
+	return io.MultiReader(headerBuf, td.reader)
+}
+
+// SafetensorsSize returns the total size of the safetensors-wrapped tensor.
+func (td *TensorData) SafetensorsSize() int64 {
+	headerJSON := td.safetensorsHeader()
+	return 8 + int64(len(headerJSON)) + td.Size
+}
+
+// NewTensorDataFromBytes creates a TensorData from raw tensor bytes.
+// This is useful for constructing packed blobs from already-extracted data.
+func NewTensorDataFromBytes(name, dtype string, shape []int32, rawData []byte) *TensorData {
+	return &TensorData{
+		Name:   name,
+		Dtype:  dtype,
+		Shape:  shape,
+		Size:   int64(len(rawData)),
+		reader: io.NewSectionReader(bytes.NewReader(rawData), 0, int64(len(rawData))),
+	}
+}
+
+// NewTensorDataFromReaderAt creates a TensorData backed by an arbitrary
+// io.ReaderAt. This is useful for constructing large synthetic tensors from
+// temporary files without loading the full payload into memory.
+func NewTensorDataFromReaderAt(name, dtype string, shape []int32, readerAt io.ReaderAt, size int64) *TensorData {
+	return &TensorData{
+		Name:   name,
+		Dtype:  dtype,
+		Shape:  shape,
+		Size:   size,
+		reader: io.NewSectionReader(readerAt, 0, size),
+	}
+}
+
+// ExtractRawFromSafetensors reads a safetensors-wrapped reader and extracts
+// the raw tensor data bytes (stripping the header).
+func ExtractRawFromSafetensors(r io.Reader) ([]byte, error) {
+	// Read header size (8 bytes, little endian)
+	var headerSize uint64
+	if err := binary.Read(r, binary.LittleEndian, &headerSize); err != nil {
+		return nil, fmt.Errorf("failed to read header size: %w", err)
+	}
+
+	// Skip header
+	if _, err := io.CopyN(io.Discard, r, int64(headerSize)); err != nil {
+		return nil, fmt.Errorf("failed to skip header: %w", err)
+	}
+
+	// Read remaining bytes (the raw tensor data)
+	return io.ReadAll(r)
+}
+
+// BuildPackedSafetensorsReader builds a streaming io.Reader that outputs a valid
+// safetensors file containing multiple tensors. Used for packing expert tensors
+// into a single blob without loading all data into memory.
+// Each TensorData must have been obtained from GetTensor.
+func BuildPackedSafetensorsReader(tensors []*TensorData) io.Reader {
+	return BuildPackedSafetensorsReaderWithMetadata(tensors, nil)
+}
+
+// BuildPackedSafetensorsReaderWithMetadata builds a streaming io.Reader that
+// outputs a valid safetensors file containing multiple tensors and optional
+// metadata.
+func BuildPackedSafetensorsReaderWithMetadata(tensors []*TensorData, metadata map[string]string) io.Reader {
+	// Build the header with sequential data offsets
+	header := make(map[string]any, len(tensors)+1)
+	var offset int
+	for _, td := range tensors {
+		header[td.Name] = tensorInfo{
+			Dtype:       td.Dtype,
+			Shape:       td.Shape,
+			DataOffsets: [2]int{offset, offset + int(td.Size)},
+		}
+		offset += int(td.Size)
+	}
+	if len(metadata) > 0 {
+		header["__metadata__"] = metadata
+	}
+
+	headerJSON, _ := json.Marshal(header)
+
+	// Pad header to 8-byte alignment
+	padding := (8 - len(headerJSON)%8) % 8
+	headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...)
+
+	// Build header with size prefix
+	headerBuf := new(bytes.Buffer)
+	binary.Write(headerBuf, binary.LittleEndian, uint64(len(headerJSON)))
+	headerBuf.Write(headerJSON)
+
+	// Build multi-reader: header + all tensor data readers
+	readers := make([]io.Reader, 0, 1+len(tensors))
+	readers = append(readers, headerBuf)
+	for _, td := range tensors {
+		td.reader.Seek(0, io.SeekStart)
+		readers = append(readers, td.reader)
+	}
+
+	return io.MultiReader(readers...)
+}
+
+// OpenForExtraction opens a safetensors file for tensor extraction.
+// The caller must call Close() when done.
+func OpenForExtraction(path string) (*TensorExtractor, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file: %w", err)
+	}
+
+	var headerSize uint64
+	if err := binary.Read(f, binary.LittleEndian, &headerSize); err != nil {
+		f.Close()
+		return nil, fmt.Errorf("failed to read header size: %w", err)
+	}
+
+	headerBytes := make([]byte, headerSize)
+	if _, err := f.Read(headerBytes); err != nil {
+		f.Close()
+		return nil, fmt.Errorf("failed to read header: %w", err)
+	}
+
+	var header map[string]tensorInfo
+	if err := json.Unmarshal(headerBytes, &header); err != nil {
+		f.Close()
+		return nil, fmt.Errorf("failed to parse header: %w", err)
+	}
+
+	delete(header, "__metadata__")
+
+	return &TensorExtractor{
+		file:       f,
+		dataOffset: 8 + int64(headerSize), // 8 bytes for header size + header content
+		header:     header,
+	}, nil
+}
+
+// GetTensor returns tensor metadata and a reader for extracting a single tensor.
+func (te *TensorExtractor) GetTensor(name string) (*TensorData, error) {
+	info, ok := te.header[name]
+	if !ok {
+		return nil, fmt.Errorf("tensor %q not found", name)
+	}
+
+	start := te.dataOffset + int64(info.DataOffsets[0])
+	size := int64(info.DataOffsets[1] - info.DataOffsets[0])
+
+	return &TensorData{
+		Name:   name,
+		Dtype:  info.Dtype,
+		Shape:  info.Shape,
+		Size:   size,
+		reader: io.NewSectionReader(te.file, start, size),
+	}, nil
+}
+
+// ListTensors returns all tensor names in sorted order.
+func (te *TensorExtractor) ListTensors() []string {
+	names := make([]string, 0, len(te.header))
+	for name := range te.header {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	return names
+}
+
+// TensorCount returns the number of tensors in the file.
+func (te *TensorExtractor) TensorCount() int {
+	return len(te.header)
+}
+
+// Close closes the underlying file.
+func (te *TensorExtractor) Close() error {
+	return te.file.Close()
+}
+
+// ExtractAll returns TensorData for all tensors in the file.
+// Each TensorData has a reader that reads from the original file.
+// The caller must call Close() on the TensorExtractor when done.
+func (te *TensorExtractor) ExtractAll() ([]*TensorData, error) {
+	names := te.ListTensors()
+	tensors := make([]*TensorData, 0, len(names))
+
+	for _, name := range names {
+		td, err := te.GetTensor(name)
+		if err != nil {
+			return nil, err
+		}
+		tensors = append(tensors, td)
+	}
+
+	return tensors, nil
+}

x/safetensors/extractor_test.go

@@ -0,0 +1,394 @@
+package safetensors
+
+import (
+	"bytes"
+	"encoding/binary"
+	"encoding/json"
+	"io"
+	"os"
+	"path/filepath"
+	"slices"
+	"testing"
+)
+
+// createTestSafetensors creates a minimal valid safetensors file with the given tensors.
+func createTestSafetensors(t *testing.T, path string, tensors map[string]struct {
+	dtype string
+	shape []int32
+	data  []byte
+},
+) {
+	t.Helper()
+
+	header := make(map[string]tensorInfo)
+	var offset int
+	var allData []byte
+
+	// Sort names for deterministic file layout
+	names := make([]string, 0, len(tensors))
+	for name := range tensors {
+		names = append(names, name)
+	}
+	slices.Sort(names)
+
+	for _, name := range names {
+		info := tensors[name]
+		header[name] = tensorInfo{
+			Dtype:       info.dtype,
+			Shape:       info.shape,
+			DataOffsets: [2]int{offset, offset + len(info.data)},
+		}
+		allData = append(allData, info.data...)
+		offset += len(info.data)
+	}
+
+	headerJSON, err := json.Marshal(header)
+	if err != nil {
+		t.Fatalf("failed to marshal header: %v", err)
+	}
+
+	// Pad to 8-byte alignment
+	padding := (8 - len(headerJSON)%8) % 8
+	headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...)
+
+	f, err := os.Create(path)
+	if err != nil {
+		t.Fatalf("failed to create file: %v", err)
+	}
+	defer f.Close()
+
+	if err := binary.Write(f, binary.LittleEndian, uint64(len(headerJSON))); err != nil {
+		t.Fatalf("failed to write header size: %v", err)
+	}
+	if _, err := f.Write(headerJSON); err != nil {
+		t.Fatalf("failed to write header: %v", err)
+	}
+	if _, err := f.Write(allData); err != nil {
+		t.Fatalf("failed to write data: %v", err)
+	}
+}
+
+func TestOpenForExtraction(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	// 4 float32 values = 16 bytes
+	data := make([]byte, 16)
+	binary.LittleEndian.PutUint32(data[0:4], 0x3f800000)   // 1.0
+	binary.LittleEndian.PutUint32(data[4:8], 0x40000000)   // 2.0
+	binary.LittleEndian.PutUint32(data[8:12], 0x40400000)  // 3.0
+	binary.LittleEndian.PutUint32(data[12:16], 0x40800000) // 4.0
+
+	createTestSafetensors(t, path, map[string]struct {
+		dtype string
+		shape []int32
+		data  []byte
+	}{
+		"test_tensor": {dtype: "F32", shape: []int32{2, 2}, data: data},
+	})
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	if ext.TensorCount() != 1 {
+		t.Errorf("TensorCount() = %d, want 1", ext.TensorCount())
+	}
+
+	names := ext.ListTensors()
+	if len(names) != 1 || names[0] != "test_tensor" {
+		t.Errorf("ListTensors() = %v, want [test_tensor]", names)
+	}
+}
+
+func TestGetTensor(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	data := make([]byte, 16)
+	for i := range 4 {
+		binary.LittleEndian.PutUint32(data[i*4:], uint32(i+1))
+	}
+
+	createTestSafetensors(t, path, map[string]struct {
+		dtype string
+		shape []int32
+		data  []byte
+	}{
+		"weight": {dtype: "F32", shape: []int32{2, 2}, data: data},
+	})
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	td, err := ext.GetTensor("weight")
+	if err != nil {
+		t.Fatalf("GetTensor failed: %v", err)
+	}
+
+	if td.Name != "weight" {
+		t.Errorf("Name = %q, want %q", td.Name, "weight")
+	}
+	if td.Dtype != "F32" {
+		t.Errorf("Dtype = %q, want %q", td.Dtype, "F32")
+	}
+	if td.Size != 16 {
+		t.Errorf("Size = %d, want 16", td.Size)
+	}
+	if len(td.Shape) != 2 || td.Shape[0] != 2 || td.Shape[1] != 2 {
+		t.Errorf("Shape = %v, want [2 2]", td.Shape)
+	}
+
+	// Read the raw data
+	rawData, err := io.ReadAll(td.Reader())
+	if err != nil {
+		t.Fatalf("Reader() read failed: %v", err)
+	}
+	if len(rawData) != 16 {
+		t.Errorf("raw data length = %d, want 16", len(rawData))
+	}
+}
+
+func TestGetTensor_NotFound(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	createTestSafetensors(t, path, map[string]struct {
+		dtype string
+		shape []int32
+		data  []byte
+	}{
+		"exists": {dtype: "F32", shape: []int32{1}, data: make([]byte, 4)},
+	})
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	_, err = ext.GetTensor("missing")
+	if err == nil {
+		t.Error("expected error for missing tensor, got nil")
+	}
+}
+
+func TestSafetensorsReaderRoundTrip(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	data := make([]byte, 16)
+	for i := range 4 {
+		binary.LittleEndian.PutUint32(data[i*4:], uint32(0x3f800000+i))
+	}
+
+	createTestSafetensors(t, path, map[string]struct {
+		dtype string
+		shape []int32
+		data  []byte
+	}{
+		"tensor_a": {dtype: "F32", shape: []int32{2, 2}, data: data},
+	})
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	td, err := ext.GetTensor("tensor_a")
+	if err != nil {
+		t.Fatalf("GetTensor failed: %v", err)
+	}
+
+	// Wrap as safetensors and extract back
+	stReader := td.SafetensorsReader()
+	stData, err := io.ReadAll(stReader)
+	if err != nil {
+		t.Fatalf("SafetensorsReader read failed: %v", err)
+	}
+
+	// Verify size
+	if int64(len(stData)) != td.SafetensorsSize() {
+		t.Errorf("SafetensorsSize() = %d, actual = %d", td.SafetensorsSize(), len(stData))
+	}
+
+	// Extract raw data back
+	raw, err := ExtractRawFromSafetensors(bytes.NewReader(stData))
+	if err != nil {
+		t.Fatalf("ExtractRawFromSafetensors failed: %v", err)
+	}
+
+	if !bytes.Equal(raw, data) {
+		t.Errorf("round-trip data mismatch: got %v, want %v", raw, data)
+	}
+}
+
+func TestNewTensorDataFromBytes(t *testing.T) {
+	data := []byte{1, 2, 3, 4}
+	td := NewTensorDataFromBytes("test", "U8", []int32{4}, data)
+
+	if td.Name != "test" {
+		t.Errorf("Name = %q, want %q", td.Name, "test")
+	}
+	if td.Size != 4 {
+		t.Errorf("Size = %d, want 4", td.Size)
+	}
+
+	rawData, err := io.ReadAll(td.Reader())
+	if err != nil {
+		t.Fatalf("Reader() failed: %v", err)
+	}
+	if !bytes.Equal(rawData, data) {
+		t.Errorf("data mismatch: got %v, want %v", rawData, data)
+	}
+}
+
+func TestBuildPackedSafetensorsReader(t *testing.T) {
+	data1 := []byte{1, 2, 3, 4}
+	data2 := []byte{5, 6, 7, 8, 9, 10, 11, 12}
+
+	td1 := NewTensorDataFromBytes("a", "U8", []int32{4}, data1)
+	td2 := NewTensorDataFromBytes("b", "U8", []int32{8}, data2)
+
+	packed := BuildPackedSafetensorsReader([]*TensorData{td1, td2})
+	packedBytes, err := io.ReadAll(packed)
+	if err != nil {
+		t.Fatalf("BuildPackedSafetensorsReader read failed: %v", err)
+	}
+
+	// Verify it's a valid safetensors file by parsing the header
+	var headerSize uint64
+	if err := binary.Read(bytes.NewReader(packedBytes), binary.LittleEndian, &headerSize); err != nil {
+		t.Fatalf("failed to read header size: %v", err)
+	}
+
+	headerJSON := packedBytes[8 : 8+headerSize]
+	var header map[string]tensorInfo
+	if err := json.Unmarshal(headerJSON, &header); err != nil {
+		t.Fatalf("failed to parse header: %v", err)
+	}
+
+	if len(header) != 2 {
+		t.Errorf("header has %d entries, want 2", len(header))
+	}
+
+	infoA, ok := header["a"]
+	if !ok {
+		t.Fatal("tensor 'a' not found in header")
+	}
+	if infoA.Dtype != "U8" {
+		t.Errorf("tensor 'a' dtype = %q, want %q", infoA.Dtype, "U8")
+	}
+
+	infoB, ok := header["b"]
+	if !ok {
+		t.Fatal("tensor 'b' not found in header")
+	}
+
+	// Verify data region contains both tensors
+	dataStart := 8 + int(headerSize)
+	dataRegion := packedBytes[dataStart:]
+	if infoA.DataOffsets[0] == 0 {
+		// a comes first
+		if !bytes.Equal(dataRegion[:4], data1) {
+			t.Error("tensor 'a' data mismatch")
+		}
+		if !bytes.Equal(dataRegion[infoB.DataOffsets[0]:infoB.DataOffsets[1]], data2) {
+			t.Error("tensor 'b' data mismatch")
+		}
+	} else {
+		// b comes first
+		if !bytes.Equal(dataRegion[:8], data2) {
+			t.Error("tensor 'b' data mismatch")
+		}
+	}
+}
+
+func TestExtractAll(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	createTestSafetensors(t, path, map[string]struct {
+		dtype string
+		shape []int32
+		data  []byte
+	}{
+		"alpha": {dtype: "F32", shape: []int32{2}, data: make([]byte, 8)},
+		"beta":  {dtype: "F16", shape: []int32{4}, data: make([]byte, 8)},
+	})
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	tensors, err := ext.ExtractAll()
+	if err != nil {
+		t.Fatalf("ExtractAll failed: %v", err)
+	}
+
+	if len(tensors) != 2 {
+		t.Errorf("ExtractAll returned %d tensors, want 2", len(tensors))
+	}
+
+	// Verify sorted order
+	if tensors[0].Name != "alpha" || tensors[1].Name != "beta" {
+		t.Errorf("tensors not in sorted order: %s, %s", tensors[0].Name, tensors[1].Name)
+	}
+}
+
+func TestExtractRawFromSafetensors_InvalidInput(t *testing.T) {
+	// Empty reader
+	_, err := ExtractRawFromSafetensors(bytes.NewReader(nil))
+	if err == nil {
+		t.Error("expected error for empty reader")
+	}
+
+	// Truncated header size
+	_, err = ExtractRawFromSafetensors(bytes.NewReader([]byte{1, 2, 3}))
+	if err == nil {
+		t.Error("expected error for truncated header size")
+	}
+}
+
+func TestOpenForExtraction_MetadataIgnored(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test.safetensors")
+
+	// Manually create a safetensors file with __metadata__
+	header := map[string]any{
+		"__metadata__": map[string]string{"format": "pt"},
+		"weight": tensorInfo{
+			Dtype:       "F32",
+			Shape:       []int32{2},
+			DataOffsets: [2]int{0, 8},
+		},
+	}
+	headerJSON, _ := json.Marshal(header)
+	padding := (8 - len(headerJSON)%8) % 8
+	headerJSON = append(headerJSON, bytes.Repeat([]byte(" "), padding)...)
+
+	f, _ := os.Create(path)
+	binary.Write(f, binary.LittleEndian, uint64(len(headerJSON)))
+	f.Write(headerJSON)
+	f.Write(make([]byte, 8))
+	f.Close()
+
+	ext, err := OpenForExtraction(path)
+	if err != nil {
+		t.Fatalf("OpenForExtraction failed: %v", err)
+	}
+	defer ext.Close()
+
+	// __metadata__ should be stripped
+	if ext.TensorCount() != 1 {
+		t.Errorf("TensorCount() = %d, want 1 (metadata should be stripped)", ext.TensorCount())
+	}
+}