mirror of https://gitee.com/namelin2022/ollama
Browse Source
With support for multimodal models becoming more varied and common it is important for clients to be able to easily see what capabilities a model has. Retuning these from the show endpoint will allow clients to easily see what a model can do.mxyng/v3
Bruce MacDonald
1 year ago
committed by
GitHub
GitHub
9 changed files with 521 additions and 69 deletions
@ -0,0 +1,360 @@ |
|||
package server |
|||
|
|||
import ( |
|||
"bytes" |
|||
"encoding/binary" |
|||
"os" |
|||
"path/filepath" |
|||
"strings" |
|||
"testing" |
|||
|
|||
"github.com/ollama/ollama/template" |
|||
"github.com/ollama/ollama/types/model" |
|||
) |
|||
|
|||
// Constants for GGUF magic bytes and version
|
|||
var ( |
|||
ggufMagic = []byte{0x47, 0x47, 0x55, 0x46} // "GGUF"
|
|||
ggufVer = uint32(3) // Version 3
|
|||
) |
|||
|
|||
// Helper function to create mock GGUF data
|
|||
func createMockGGUFData(architecture string, vision bool) []byte { |
|||
var buf bytes.Buffer |
|||
|
|||
// Write GGUF header
|
|||
buf.Write(ggufMagic) |
|||
binary.Write(&buf, binary.LittleEndian, ggufVer) |
|||
|
|||
// Write tensor count (0 for our test)
|
|||
var numTensors uint64 = 0 |
|||
binary.Write(&buf, binary.LittleEndian, numTensors) |
|||
|
|||
// Calculate number of metadata entries
|
|||
numMetaEntries := uint64(1) // architecture entry
|
|||
if vision { |
|||
numMetaEntries++ |
|||
} |
|||
// Add embedding entry if architecture is "bert"
|
|||
if architecture == "bert" { |
|||
numMetaEntries++ |
|||
} |
|||
binary.Write(&buf, binary.LittleEndian, numMetaEntries) |
|||
|
|||
// Write architecture metadata
|
|||
archKey := "general.architecture" |
|||
keyLen := uint64(len(archKey)) |
|||
binary.Write(&buf, binary.LittleEndian, keyLen) |
|||
buf.WriteString(archKey) |
|||
|
|||
// String type (8)
|
|||
var strType uint32 = 8 |
|||
binary.Write(&buf, binary.LittleEndian, strType) |
|||
|
|||
// String length
|
|||
strLen := uint64(len(architecture)) |
|||
binary.Write(&buf, binary.LittleEndian, strLen) |
|||
buf.WriteString(architecture) |
|||
|
|||
if vision { |
|||
visionKey := architecture + ".vision.block_count" |
|||
keyLen = uint64(len(visionKey)) |
|||
binary.Write(&buf, binary.LittleEndian, keyLen) |
|||
buf.WriteString(visionKey) |
|||
|
|||
// uint32 type (4)
|
|||
var uint32Type uint32 = 4 |
|||
binary.Write(&buf, binary.LittleEndian, uint32Type) |
|||
|
|||
// uint32 value (1)
|
|||
var countVal uint32 = 1 |
|||
binary.Write(&buf, binary.LittleEndian, countVal) |
|||
} |
|||
// Write embedding metadata if architecture is "bert"
|
|||
if architecture == "bert" { |
|||
poolKey := architecture + ".pooling_type" |
|||
keyLen = uint64(len(poolKey)) |
|||
binary.Write(&buf, binary.LittleEndian, keyLen) |
|||
buf.WriteString(poolKey) |
|||
|
|||
// uint32 type (4)
|
|||
var uint32Type uint32 = 4 |
|||
binary.Write(&buf, binary.LittleEndian, uint32Type) |
|||
|
|||
// uint32 value (1)
|
|||
var poolingVal uint32 = 1 |
|||
binary.Write(&buf, binary.LittleEndian, poolingVal) |
|||
} |
|||
|
|||
return buf.Bytes() |
|||
} |
|||
|
|||
func TestModelCapabilities(t *testing.T) { |
|||
// Create a temporary directory for test files
|
|||
tempDir, err := os.MkdirTemp("", "model_capabilities_test") |
|||
if err != nil { |
|||
t.Fatalf("Failed to create temp directory: %v", err) |
|||
} |
|||
defer os.RemoveAll(tempDir) |
|||
|
|||
// Create different types of mock model files
|
|||
completionModelPath := filepath.Join(tempDir, "model.bin") |
|||
visionModelPath := filepath.Join(tempDir, "vision_model.bin") |
|||
embeddingModelPath := filepath.Join(tempDir, "embedding_model.bin") |
|||
// Create a simple model file for tests that don't depend on GGUF content
|
|||
simpleModelPath := filepath.Join(tempDir, "simple_model.bin") |
|||
|
|||
err = os.WriteFile(completionModelPath, createMockGGUFData("llama", false), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create completion model file: %v", err) |
|||
} |
|||
err = os.WriteFile(visionModelPath, createMockGGUFData("llama", true), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create completion model file: %v", err) |
|||
} |
|||
err = os.WriteFile(embeddingModelPath, createMockGGUFData("bert", false), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create embedding model file: %v", err) |
|||
} |
|||
err = os.WriteFile(simpleModelPath, []byte("dummy model data"), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create simple model file: %v", err) |
|||
} |
|||
|
|||
toolsInsertTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
chatTemplate, err := template.Parse("{{ .prompt }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
toolsTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
|
|||
testModels := []struct { |
|||
name string |
|||
model Model |
|||
expectedCaps []model.Capability |
|||
}{ |
|||
{ |
|||
name: "model with completion capability", |
|||
model: Model{ |
|||
ModelPath: completionModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityCompletion}, |
|||
}, |
|||
|
|||
{ |
|||
name: "model with completion, tools, and insert capability", |
|||
model: Model{ |
|||
ModelPath: completionModelPath, |
|||
Template: toolsInsertTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityTools, model.CapabilityInsert}, |
|||
}, |
|||
{ |
|||
name: "model with tools and insert capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: toolsInsertTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityTools, model.CapabilityInsert}, |
|||
}, |
|||
{ |
|||
name: "model with tools capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: toolsTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityTools}, |
|||
}, |
|||
{ |
|||
name: "model with vision capability", |
|||
model: Model{ |
|||
ModelPath: visionModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision}, |
|||
}, |
|||
{ |
|||
name: "model with vision, tools, and insert capability", |
|||
model: Model{ |
|||
ModelPath: visionModelPath, |
|||
Template: toolsInsertTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityCompletion, model.CapabilityVision, model.CapabilityTools, model.CapabilityInsert}, |
|||
}, |
|||
{ |
|||
name: "model with embedding capability", |
|||
model: Model{ |
|||
ModelPath: embeddingModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
expectedCaps: []model.Capability{model.CapabilityEmbedding}, |
|||
}, |
|||
} |
|||
|
|||
// compare two slices of model.Capability regardless of order
|
|||
compareCapabilities := func(a, b []model.Capability) bool { |
|||
if len(a) != len(b) { |
|||
return false |
|||
} |
|||
|
|||
aCount := make(map[model.Capability]int) |
|||
for _, cap := range a { |
|||
aCount[cap]++ |
|||
} |
|||
|
|||
bCount := make(map[model.Capability]int) |
|||
for _, cap := range b { |
|||
bCount[cap]++ |
|||
} |
|||
|
|||
for cap, count := range aCount { |
|||
if bCount[cap] != count { |
|||
return false |
|||
} |
|||
} |
|||
|
|||
return true |
|||
} |
|||
|
|||
for _, tt := range testModels { |
|||
t.Run(tt.name, func(t *testing.T) { |
|||
// Test Capabilities method
|
|||
caps := tt.model.Capabilities() |
|||
if !compareCapabilities(caps, tt.expectedCaps) { |
|||
t.Errorf("Expected capabilities %v, got %v", tt.expectedCaps, caps) |
|||
} |
|||
}) |
|||
} |
|||
} |
|||
|
|||
func TestModelCheckCapabilities(t *testing.T) { |
|||
// Create a temporary directory for test files
|
|||
tempDir, err := os.MkdirTemp("", "model_check_capabilities_test") |
|||
if err != nil { |
|||
t.Fatalf("Failed to create temp directory: %v", err) |
|||
} |
|||
defer os.RemoveAll(tempDir) |
|||
|
|||
visionModelPath := filepath.Join(tempDir, "vision_model.bin") |
|||
simpleModelPath := filepath.Join(tempDir, "model.bin") |
|||
embeddingModelPath := filepath.Join(tempDir, "embedding_model.bin") |
|||
|
|||
err = os.WriteFile(simpleModelPath, []byte("dummy model data"), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create simple model file: %v", err) |
|||
} |
|||
err = os.WriteFile(visionModelPath, createMockGGUFData("llama", true), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create vision model file: %v", err) |
|||
} |
|||
err = os.WriteFile(embeddingModelPath, createMockGGUFData("bert", false), 0o644) |
|||
if err != nil { |
|||
t.Fatalf("Failed to create embedding model file: %v", err) |
|||
} |
|||
|
|||
toolsInsertTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}{{ if .suffix }}{{ .suffix }}{{ end }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
chatTemplate, err := template.Parse("{{ .prompt }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
toolsTemplate, err := template.Parse("{{ .prompt }}{{ if .tools }}{{ .tools }}{{ end }}") |
|||
if err != nil { |
|||
t.Fatalf("Failed to parse template: %v", err) |
|||
} |
|||
|
|||
tests := []struct { |
|||
name string |
|||
model Model |
|||
checkCaps []model.Capability |
|||
expectedErrMsg string |
|||
}{ |
|||
{ |
|||
name: "completion model without tools capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityTools}, |
|||
expectedErrMsg: "does not support tools", |
|||
}, |
|||
{ |
|||
name: "model with all needed capabilities", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: toolsInsertTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityTools, model.CapabilityInsert}, |
|||
}, |
|||
{ |
|||
name: "model missing insert capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: toolsTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityInsert}, |
|||
expectedErrMsg: "does not support insert", |
|||
}, |
|||
{ |
|||
name: "model missing vision capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: toolsTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityVision}, |
|||
expectedErrMsg: "does not support vision", |
|||
}, |
|||
{ |
|||
name: "model with vision capability", |
|||
model: Model{ |
|||
ModelPath: visionModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityVision}, |
|||
}, |
|||
{ |
|||
name: "model with embedding capability", |
|||
model: Model{ |
|||
ModelPath: embeddingModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{model.CapabilityEmbedding}, |
|||
}, |
|||
{ |
|||
name: "unknown capability", |
|||
model: Model{ |
|||
ModelPath: simpleModelPath, |
|||
Template: chatTemplate, |
|||
}, |
|||
checkCaps: []model.Capability{"unknown"}, |
|||
expectedErrMsg: "unknown capability", |
|||
}, |
|||
} |
|||
|
|||
for _, tt := range tests { |
|||
t.Run(tt.name, func(t *testing.T) { |
|||
// Test CheckCapabilities method
|
|||
err := tt.model.CheckCapabilities(tt.checkCaps...) |
|||
if tt.expectedErrMsg == "" { |
|||
if err != nil { |
|||
t.Errorf("Expected no error, got: %v", err) |
|||
} |
|||
} else { |
|||
if err == nil { |
|||
t.Errorf("Expected error containing %q, got nil", tt.expectedErrMsg) |
|||
} else if !strings.Contains(err.Error(), tt.expectedErrMsg) { |
|||
t.Errorf("Expected error containing %q, got: %v", tt.expectedErrMsg, err) |
|||
} |
|||
} |
|||
}) |
|||
} |
|||
} |
|||
@ -0,0 +1,15 @@ |
|||
package model |
|||
|
|||
type Capability string |
|||
|
|||
const ( |
|||
CapabilityCompletion = Capability("completion") |
|||
CapabilityTools = Capability("tools") |
|||
CapabilityInsert = Capability("insert") |
|||
CapabilityVision = Capability("vision") |
|||
CapabilityEmbedding = Capability("embedding") |
|||
) |
|||
|
|||
func (c Capability) String() string { |
|||
return string(c) |
|||
} |
|||
Loading…
Reference in new issue