diff --git a/x/create/client/create.go b/x/create/client/create.go index d8a1a8f80..a6f1cd0f4 100644 --- a/x/create/client/create.go +++ b/x/create/client/create.go @@ -243,6 +243,40 @@ func appendLayersManifestWriter(next create.ManifestWriter, extra []create.Layer } } +func draftMetadata(draftDir string) (*model.Draft, error) { + configPath := filepath.Join(draftDir, "config.json") + data, err := os.ReadFile(configPath) + if err != nil { + return nil, fmt.Errorf("failed to read draft config %s: %w", configPath, err) + } + + var cfg struct { + Architectures []string `json:"architectures"` + ModelType string `json:"model_type"` + } + if err := json.Unmarshal(data, &cfg); err != nil { + return nil, fmt.Errorf("failed to parse draft config %s: %w", configPath, err) + } + + arch := "" + if len(cfg.Architectures) > 0 { + arch = cfg.Architectures[0] + } + if arch == "" { + arch = cfg.ModelType + } + if arch == "" { + return nil, fmt.Errorf("draft architecture not found in %s", configPath) + } + + return &model.Draft{ + ModelFormat: "safetensors", + Architecture: arch, + TensorPrefix: "draft.", + Config: "draft/config.json", + }, nil +} + func createModelFromBaseWithDraft(opts CreateOptions, draftLayers []create.LayerInfo, progressFn func(string)) error { progressFn(fmt.Sprintf("loading base model %s", opts.ModelDir)) baseManifest, err := imagemanifest.LoadManifest(opts.ModelDir) @@ -487,12 +521,11 @@ func newManifestWriter(opts CreateOptions, capabilities []string, parserName, re configData.Parser = resolveParserName(opts.Modelfile, parserName) configData.Renderer = resolveRendererName(opts.Modelfile, rendererName) if opts.Modelfile != nil && opts.Modelfile.Draft != "" { - configData.Draft = &model.Draft{ - ModelFormat: "safetensors", - Architecture: "Gemma4AssistantForCausalLM", - TensorPrefix: "draft.", - Config: "draft/config.json", + draft, err := draftMetadata(opts.Modelfile.Draft) + if err != nil { + return err } + configData.Draft = draft } configJSON, err := json.Marshal(configData) if err != nil { diff --git a/x/create/client/create_test.go b/x/create/client/create_test.go index bc8f4d6e9..47cd9d8b3 100644 --- a/x/create/client/create_test.go +++ b/x/create/client/create_test.go @@ -544,10 +544,15 @@ func TestNewManifestWriter_PopulatesFileTypeFromQuantize(t *testing.T) { func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { t.Setenv("OLLAMA_MODELS", t.TempDir()) + draftDir := t.TempDir() + if err := os.WriteFile(filepath.Join(draftDir, "config.json"), []byte(`{"architectures":["DFlashDraftModel"],"model_type":"qwen3"}`), 0o644); err != nil { + t.Fatalf("WriteFile() error = %v", err) + } + opts := CreateOptions{ ModelName: "test-draft", ModelDir: t.TempDir(), - Modelfile: &ModelfileConfig{Draft: "/tmp/assistant"}, + Modelfile: &ModelfileConfig{Draft: draftDir}, } writer := newManifestWriter(opts, []string{"completion"}, "gemma4", "gemma4") @@ -581,6 +586,9 @@ func TestNewManifestWriter_PopulatesDraftMetadata(t *testing.T) { if cfg.Draft.TensorPrefix != "draft." || cfg.Draft.Config != "draft/config.json" { t.Fatalf("Draft = %#v, want draft prefix/config", cfg.Draft) } + if cfg.Draft.Architecture != "DFlashDraftModel" { + t.Fatalf("Draft architecture = %q, want DFlashDraftModel", cfg.Draft.Architecture) + } } func TestSupportsThinking(t *testing.T) {