Repository: ollama/ollama
File: server/internal/cmd/opp/internal/safetensors/safetensors.go (Go)

Code before:
		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
	}

	endOfHeader := 8 + headerSize // 8 bytes for header size plus the header itself

	// 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.") {
			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 := endOfHeader + v.Offsets[0]
		end := endOfHeader + 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,

Review comment (suggestion):
This will miss the following (e.g. model-00004-of-00004.safetensors in Meta-Llama-3.1-8B-Instruct
```json
  "lm_head.weight": {
    "dtype": "BF16",
    "shape": [
      128256,
      4096
    ],
    "data_offsets": [
      0,
      1050673152
    ]
  },
```

Code after:
		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
	}

	endOfHeader := 8 + headerSize // 8 bytes for header size plus the header itself

	// 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 name == "__metadata__" {
			// TODO(bmizerany): do something with metadata?
			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 := endOfHeader + v.Offsets[0]
		end := endOfHeader + 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{