Optionally set context-size in artifact (docker#92)

Signed-off-by: Emily Casey <emily.casey@docker.com>

Author: ekcasey

pkg/distribution/README.md

@@ -31,6 +31,9 @@ make build
 # Package a model with license files and push to a registry
 ./bin/model-distribution-tool package --licenses license1.txt --licenses license2.txt ./model.gguf registry.example.com/models/llama:v1.0
 
+# Package a model with a default context size and push to a registry
+./bin/model-distribution-tool ./model.gguf --context-size 2048 registry.example.com/models/llama:v1.0
+
 # Push a model from the content store to the registry
 ./bin/model-distribution-tool push registry.example.com/models/llama:v1.0
 

pkg/distribution/builder/builder.go

@@ -38,6 +38,12 @@ func (b *Builder) WithLicense(path string) (*Builder, error) {
 	}, nil
 }
 
+func (b *Builder) WithContextSize(size uint64) *Builder {
+	return &Builder{
+		model: mutate.ContextSize(b.model, size),
+	}
+}
+
 // Target represents a build target
 type Target interface {
 	Write(context.Context, types.ModelArtifact, io.Writer) error

pkg/distribution/internal/mutate/model.go

@@ -16,6 +16,7 @@ type model struct {
 	base            types.ModelArtifact
 	appended        []v1.Layer
 	configMediaType ggcr.MediaType
+	contextSize     *uint64
 }
 
 func (m *model) Descriptor() (types.Descriptor, error) {
@@ -123,6 +124,9 @@ func (m *model) RawConfigFile() ([]byte, error) {
 		}
 		cf.RootFS.DiffIDs = append(cf.RootFS.DiffIDs, diffID)
 	}
+	if m.contextSize != nil {
+		cf.Config.ContextSize = m.contextSize
+	}
 	raw, err := json.Marshal(cf)
 	if err != nil {
 		return nil, err

pkg/distribution/internal/mutate/mutate.go

@@ -20,3 +20,10 @@ func ConfigMediaType(mdl types.ModelArtifact, mt ggcr.MediaType) types.ModelArti
 		configMediaType: mt,
 	}
 }
+
+func ContextSize(mdl types.ModelArtifact, cs uint64) types.ModelArtifact {
+	return &model{
+		base:        mdl,
+		contextSize: &cs,
+	}
+}

pkg/distribution/internal/mutate/mutate_test.go

@@ -83,3 +83,32 @@ func TestConfigMediaTypes(t *testing.T) {
 		t.Fatalf("Expected media type %s, got %s", newMediaType, manifest2.Config.MediaType)
 	}
 }
+
+func TestContextSize(t *testing.T) {
+	mdl1, err := gguf.NewModel(filepath.Join("..", "..", "assets", "dummy.gguf"))
+	if err != nil {
+		t.Fatalf("Failed to create model: %v", err)
+	}
+	cfg, err := mdl1.Config()
+	if err != nil {
+		t.Fatalf("Failed to get config file: %v", err)
+	}
+	if cfg.ContextSize != nil {
+		t.Fatalf("Epected nil context size got %d", cfg.ContextSize)
+	}
+
+	// set the context size
+	mdl2 := mutate.ContextSize(mdl1, 2096)
+
+	// check the config
+	cfg2, err := mdl2.Config()
+	if err != nil {
+		t.Fatalf("Failed to get config file: %v", err)
+	}
+	if cfg2.ContextSize == nil {
+		t.Fatal("Expected non-nil context")
+	}
+	if *cfg2.ContextSize != uint64(2096) {
+		t.Fatalf("Expected context size of 2096 got %d", *cfg2.ContextSize)
+	}
+}

pkg/distribution/types/config.go

@@ -44,6 +44,7 @@ type Config struct {
 	Architecture string            `json:"architecture,omitempty"`
 	Size         string            `json:"size,omitempty"`
 	GGUF         map[string]string `json:"gguf,omitempty"`
+	ContextSize  *uint64           `json:"context_size,omitempty"`
 }
 
 // Descriptor provides metadata about the provenance of the model.