github.com/anchore/syft@v1.38.2/syft/pkg/cataloger/ai/test_helpers_test.go (about)

     1  package ai
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  )
     7  
     8  // GGUF type constants for test builder
     9  // https://github.com/ggml-org/ggml/blob/master/docs/gguf.md
    10  const (
    11  	ggufMagic       = 0x46554747 // "GGUF" in little-endian
    12  	ggufTypeUint8   = 0
    13  	ggufTypeInt8    = 1
    14  	ggufTypeUint16  = 2
    15  	ggufTypeInt16   = 3
    16  	ggufTypeUint32  = 4
    17  	ggufTypeInt32   = 5
    18  	ggufTypeFloat32 = 6
    19  	ggufTypeBool    = 7
    20  	ggufTypeString  = 8
    21  	ggufTypeArray   = 9
    22  	ggufTypeUint64  = 10
    23  	ggufTypeInt64   = 11
    24  	ggufTypeFloat64 = 12
    25  )
    26  
    27  // testGGUFBuilder helps build GGUF files for testing
    28  type testGGUFBuilder struct {
    29  	buf         *bytes.Buffer
    30  	version     uint32
    31  	tensorCount uint64
    32  	kvPairs     []testKVPair
    33  }
    34  
    35  type testKVPair struct {
    36  	key       string
    37  	valueType uint32
    38  	value     interface{}
    39  }
    40  
    41  func newTestGGUFBuilder() *testGGUFBuilder {
    42  	return &testGGUFBuilder{
    43  		buf:         new(bytes.Buffer),
    44  		version:     3,
    45  		tensorCount: 0,
    46  		kvPairs:     []testKVPair{},
    47  	}
    48  }
    49  
    50  func (b *testGGUFBuilder) withVersion(v uint32) *testGGUFBuilder {
    51  	b.version = v
    52  	return b
    53  }
    54  
    55  func (b *testGGUFBuilder) withTensorCount(count uint64) *testGGUFBuilder {
    56  	b.tensorCount = count
    57  	return b
    58  }
    59  
    60  func (b *testGGUFBuilder) withStringKV(key, value string) *testGGUFBuilder {
    61  	b.kvPairs = append(b.kvPairs, testKVPair{key: key, valueType: ggufTypeString, value: value})
    62  	return b
    63  }
    64  
    65  func (b *testGGUFBuilder) withUint64KV(key string, value uint64) *testGGUFBuilder {
    66  	b.kvPairs = append(b.kvPairs, testKVPair{key: key, valueType: ggufTypeUint64, value: value})
    67  	return b
    68  }
    69  
    70  func (b *testGGUFBuilder) withUint32KV(key string, value uint32) *testGGUFBuilder {
    71  	b.kvPairs = append(b.kvPairs, testKVPair{key: key, valueType: ggufTypeUint32, value: value})
    72  	return b
    73  }
    74  
    75  func (b *testGGUFBuilder) writeString(s string) {
    76  	binary.Write(b.buf, binary.LittleEndian, uint64(len(s)))
    77  	b.buf.WriteString(s)
    78  }
    79  
    80  func (b *testGGUFBuilder) build() []byte {
    81  	// Write magic number "GGUF"
    82  	binary.Write(b.buf, binary.LittleEndian, uint32(ggufMagic))
    83  
    84  	// Write version
    85  	binary.Write(b.buf, binary.LittleEndian, b.version)
    86  
    87  	// Write tensor count
    88  	binary.Write(b.buf, binary.LittleEndian, b.tensorCount)
    89  
    90  	// Write KV count
    91  	binary.Write(b.buf, binary.LittleEndian, uint64(len(b.kvPairs)))
    92  
    93  	// Write KV pairs
    94  	for _, kv := range b.kvPairs {
    95  		// Write key
    96  		b.writeString(kv.key)
    97  		// Write value type
    98  		binary.Write(b.buf, binary.LittleEndian, kv.valueType)
    99  		// Write value based on type
   100  		switch kv.valueType {
   101  		case ggufTypeString:
   102  			b.writeString(kv.value.(string))
   103  		case ggufTypeUint32:
   104  			binary.Write(b.buf, binary.LittleEndian, kv.value.(uint32))
   105  		case ggufTypeUint64:
   106  			binary.Write(b.buf, binary.LittleEndian, kv.value.(uint64))
   107  		case ggufTypeUint8:
   108  			binary.Write(b.buf, binary.LittleEndian, kv.value.(uint8))
   109  		case ggufTypeInt32:
   110  			binary.Write(b.buf, binary.LittleEndian, kv.value.(int32))
   111  		case ggufTypeBool:
   112  			var v uint8
   113  			if kv.value.(bool) {
   114  				v = 1
   115  			}
   116  			binary.Write(b.buf, binary.LittleEndian, v)
   117  		}
   118  	}
   119  
   120  	return b.buf.Bytes()
   121  }
   122  
   123  // buildInvalidMagic creates a file with invalid magic number
   124  func (b *testGGUFBuilder) buildInvalidMagic() []byte {
   125  	buf := new(bytes.Buffer)
   126  	binary.Write(buf, binary.LittleEndian, uint32(0x12345678))
   127  	return buf.Bytes()
   128  }