server/internal: copy bmizerany/ollama-go to internal package (#9294)

This commit copies (without history) the bmizerany/ollama-go repository
with the intention of integrating it into the ollama as a replacement
for the pushing, and pulling of models, and management of the cache they
are pushed and pulled from.

New homes for these packages will be determined as they are integrated
and we have a better understanding of proper package boundaries.

server/internal/cmd/opp/internal/safetensors/safetensors.go

@@ -0,0 +1,220 @@
+// safetensors provides a reader for the safetensor directories and files.
+package safetensors
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"io/fs"
+	"iter"
+	"slices"
+	"strconv"
+	"strings"
+)
+
+// Tensor represents a single tensor in a safetensors file.
+//
+// It's zero value is not valid. Use [Model.Tensors] to get valid tensors.
+//
+// It is not safe for use across multiple goroutines.
+type Tensor struct {
+	name     string
+	dataType string
+	shape    []int64
+
+	fsys   fs.FS
+	fname  string // entry name in fsys
+	offset int64
+	size   int64
+}
+
+type Model struct {
+	fsys fs.FS
+}
+
+func Read(fsys fs.FS) (*Model, error) {
+	return &Model{fsys: fsys}, nil
+}
+
+func (m *Model) Tensors() iter.Seq2[*Tensor, error] {
+	return func(yield func(*Tensor, error) bool) {
+		entries, err := fs.Glob(m.fsys, "*.safetensors")
+		if err != nil {
+			yield(nil, err)
+			return
+		}
+		for _, e := range entries {
+			tt, err := m.readTensors(e)
+			if err != nil {
+				yield(nil, err)
+				return
+			}
+			for _, t := range tt {
+				if !yield(t, nil) {
+					return
+				}
+			}
+		}
+	}
+}
+
+func (m *Model) readTensors(fname string) ([]*Tensor, error) {
+	f, err := m.fsys.Open(fname)
+	if err != nil {
+		return nil, err
+	}
+	defer f.Close()
+
+	finfo, err := f.Stat()
+	if err != nil {
+		return nil, err
+	}
+
+	headerSize, err := readInt64(f)
+	if err != nil {
+		return nil, err
+	}
+
+	data := make([]byte, headerSize)
+	_, err = io.ReadFull(f, data)
+	if err != nil {
+		return nil, err
+	}
+
+	var raws map[string]json.RawMessage
+	if err := json.Unmarshal(data, &raws); err != nil {
+		return nil, err
+	}
+
+	// TODO(bmizerany): do something with metadata? This could be another
+	// header read if needed. We also need to figure out if the metadata is
+	// present in only one .safetensors file or if each file may have their
+	// own and if it needs to follow each tensor. Currently, I (bmizerany)
+	// am only seeing them show up with one entry for file type which is
+	// always "pt".
+
+	tt := make([]*Tensor, 0, len(raws))
+	for name, raw := range raws {
+		if !strings.HasPrefix(name, "model.layer") {
+			continue
+		}
+		var v struct {
+			DataType string  `json:"dtype"`
+			Shape    []int64 `json:"shape"`
+			Offsets  []int64 `json:"data_offsets"`
+		}
+		if err := json.Unmarshal(raw, &v); err != nil {
+			return nil, fmt.Errorf("error unmarshalling layer %q: %w", name, err)
+		}
+		if len(v.Offsets) != 2 {
+			return nil, fmt.Errorf("invalid offsets for %q: %v", name, v.Offsets)
+		}
+
+		// TODO(bmizerany): after collecting, validate all offests make
+		// tensors contiguous?
+		begin, end := v.Offsets[0], v.Offsets[1]
+		if err := checkBeginEnd(finfo.Size(), begin, end); err != nil {
+			return nil, err
+		}
+
+		// TODO(bmizerany): just yield.. don't be silly and make a slice :)
+		tt = append(tt, &Tensor{
+			name:     name,
+			dataType: v.DataType,
+			shape:    v.Shape,
+			fsys:     m.fsys,
+			fname:    fname,
+			offset:   begin,
+			size:     end - begin,
+		})
+	}
+	return tt, nil
+}
+
+func checkBeginEnd(size, begin, end int64) error {
+	if begin < 0 {
+		return fmt.Errorf("begin must not be negative: %d", begin)
+	}
+	if end < 0 {
+		return fmt.Errorf("end must not be negative: %d", end)
+	}
+	if end < begin {
+		return fmt.Errorf("end must be >= begin: %d < %d", end, begin)
+	}
+	if end > size {
+		return fmt.Errorf("end must be <= size: %d > %d", end, size)
+	}
+	return nil
+}
+
+func readInt64(r io.Reader) (int64, error) {
+	var v uint64
+	var buf [8]byte
+	if _, err := io.ReadFull(r, buf[:]); err != nil {
+		return 0, err
+	}
+	for i := range buf {
+		v |= uint64(buf[i]) << (8 * i)
+	}
+	return int64(v), nil
+}
+
+type Shape []int64
+
+func (s Shape) String() string {
+	var b strings.Builder
+	b.WriteByte('[')
+	for i, v := range s {
+		if i > 0 {
+			b.WriteByte(',')
+		}
+		b.WriteString(strconv.FormatInt(v, 10))
+	}
+	b.WriteByte(']')
+	return b.String()
+}
+
+func (t *Tensor) Name() string     { return t.name }
+func (t *Tensor) DataType() string { return t.dataType }
+func (t *Tensor) Size() int64      { return t.size }
+func (t *Tensor) Shape() Shape     { return slices.Clone(t.shape) }
+
+func (t *Tensor) Reader() (io.ReadCloser, error) {
+	f, err := t.fsys.Open(t.fname)
+	if err != nil {
+		return nil, err
+	}
+	r := newSectionReader(f, t.offset, t.size)
+	rc := struct {
+		io.Reader
+		io.Closer
+	}{r, f}
+	return rc, nil
+}
+
+// newSectionReader returns a new io.Reader that reads from r starting at
+// offset. It is a convenience function for creating a io.SectionReader when r
+// may not be an io.ReaderAt.
+//
+// If r is already a ReaderAt, it is returned directly, otherwise if r is an
+// io.Seeker, a new io.ReaderAt is returned that wraps r after seeking to the
+// beginning of the file.
+//
+// If r is an io.Seeker,
+// or slow path. The slow path is used when r does not implement io.ReaderAt,
+// in which case it must discard the data it reads.
+func newSectionReader(r io.Reader, offset, n int64) io.Reader {
+	if r, ok := r.(io.ReaderAt); ok {
+		return io.NewSectionReader(r, offset, n)
+	}
+	if r, ok := r.(io.ReadSeeker); ok {
+		r.Seek(offset, io.SeekStart)
+		return io.LimitReader(r, n)
+	}
+	// Discard to offset and return a limited reader.
+	_, err := io.CopyN(io.Discard, r, offset)
+	if err != nil {
+		return nil
+	}
+	return io.LimitReader(r, n)
+}