95 lines
2.4 KiB
Go
95 lines
2.4 KiB
Go
package cache
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"os"
|
|
"slices"
|
|
"sync"
|
|
|
|
"github.com/ollama/ollama/fs/ggml"
|
|
"github.com/ollama/ollama/template"
|
|
"github.com/ollama/ollama/thinking"
|
|
"github.com/ollama/ollama/types/model"
|
|
)
|
|
|
|
// ggufCapabilities is a cache for gguf model capabilities
|
|
var ggufCapabilities = &sync.Map{}
|
|
|
|
// ModelInfo contains the minimal information needed to determine capabilities
|
|
type ModelInfo struct {
|
|
ModelPath string
|
|
ProjectorPaths []string
|
|
Template *template.Template
|
|
}
|
|
|
|
// Capabilities returns the capabilities that the model supports
|
|
func Capabilities(info ModelInfo) []model.Capability {
|
|
capabilities, err := ggufCapabilties(info.ModelPath)
|
|
if err != nil {
|
|
slog.Error("could not determine gguf capabilities", "error", err)
|
|
}
|
|
|
|
if info.Template == nil {
|
|
return capabilities
|
|
}
|
|
|
|
// Check for tools capability
|
|
if slices.Contains(info.Template.Vars(), "tools") {
|
|
capabilities = append(capabilities, model.CapabilityTools)
|
|
}
|
|
|
|
// Check for insert capability
|
|
if slices.Contains(info.Template.Vars(), "suffix") {
|
|
capabilities = append(capabilities, model.CapabilityInsert)
|
|
}
|
|
|
|
// Check for vision capability in projector-based models
|
|
if len(info.ProjectorPaths) > 0 {
|
|
capabilities = append(capabilities, model.CapabilityVision)
|
|
}
|
|
|
|
// Check for thinking capability
|
|
openingTag, closingTag := thinking.InferTags(info.Template.Template)
|
|
if openingTag != "" && closingTag != "" {
|
|
capabilities = append(capabilities, model.CapabilityThinking)
|
|
}
|
|
|
|
return capabilities
|
|
}
|
|
|
|
func ggufCapabilties(modelPath string) ([]model.Capability, error) {
|
|
if ggufCapabilities, ok := ggufCapabilities.Load(modelPath); ok {
|
|
capabilities := ggufCapabilities.([]model.Capability)
|
|
return capabilities, nil
|
|
}
|
|
|
|
// If not cached, read the model file to determine capabilities
|
|
capabilities := []model.Capability{}
|
|
|
|
r, err := os.Open(modelPath)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer r.Close()
|
|
|
|
f, err := ggml.Decode(r, 1024)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if _, ok := f.KV()[fmt.Sprintf("%s.pooling_type", f.KV().Architecture())]; ok {
|
|
capabilities = append(capabilities, model.CapabilityEmbedding)
|
|
} else {
|
|
capabilities = append(capabilities, model.CapabilityCompletion)
|
|
}
|
|
if _, ok := f.KV()[fmt.Sprintf("%s.vision.block_count", f.KV().Architecture())]; ok {
|
|
capabilities = append(capabilities, model.CapabilityVision)
|
|
}
|
|
|
|
// Cache the capabilities for future use
|
|
ggufCapabilities.Store(modelPath, capabilities)
|
|
|
|
return capabilities, nil
|
|
}
|