diff --git a/cmd/gguf-parser/README.md b/cmd/gguf-parser/README.md index 4830351..d084065 100644 --- a/cmd/gguf-parser/README.md +++ b/cmd/gguf-parser/README.md @@ -84,8 +84,9 @@ GLOBAL OPTIONS: Model/Remote/Ollama - --ol-model value Model name of Ollama, e.g. gemma2. - --ol-usage Specify respecting the extending layers introduced by Ollama, works with --ol-model, which affects the usage estimation. (default: false) + --ol-base-url value Model base URL of Ollama, e.g. https://registry.ollama.ai. (default: "https://registry.ollama.ai") + --ol-model value Model name of Ollama, e.g. gemma2. + --ol-usage Specify respecting the extending layers introduced by Ollama, works with --ol-model, which affects the usage estimation. (default: false) Output diff --git a/cmd/gguf-parser/main.go b/cmd/gguf-parser/main.go index 798090a..e06bc0e 100644 --- a/cmd/gguf-parser/main.go +++ b/cmd/gguf-parser/main.go @@ -223,6 +223,14 @@ func main() { "works with --ms-repo/--ms-file pair or --ms-draft-repo/--ms-draft-file pair. " + "See https://modelscope.cn/my/myaccesstoken.", }, + &cli.StringFlag{ + Destination: &olBaseURL, + Value: olBaseURL, + Category: "Model/Remote/Ollama", + Name: "ol-base-url", + Usage: "Model base URL of Ollama, e.g. " + + "https://registry.ollama.ai.", + }, &cli.StringFlag{ Destination: &olModel, Value: olModel, @@ -561,6 +569,7 @@ var ( msDraftRepo string // for estimate msDraftFile string // for estimate msToken string + olBaseURL = "https://registry.ollama.ai" olModel string olUsage bool // load options @@ -730,7 +739,7 @@ func mainAction(c *cli.Context) error { } gf, err = ParseGGUFFileFromModelScope(ctx, msRepo, msFile, ropts...) case olModel != "": - om := ParseOllamaModel(olModel) + om := ParseOllamaModel(olModel, SetOllamaModelBaseURL(olBaseURL)) gf, err = ParseGGUFFileFromOllamaModel(ctx, om, ropts...) if om != nil && olUsage { // Parameters override. diff --git a/ollama_model.go b/ollama_model.go index c64fd63..3ffd950 100644 --- a/ollama_model.go +++ b/ollama_model.go @@ -69,17 +69,36 @@ type ( // ParseOllamaModel parses the given Ollama model string, // and returns the OllamaModel, or nil if the model is invalid. -func ParseOllamaModel(model string) *OllamaModel { +func ParseOllamaModel(model string, opts ...OllamaModelOption) *OllamaModel { if model == "" { return nil } + var o _OllamaModelOptions + for _, opt := range opts { + opt(&o) + } + om := OllamaModel{ Schema: OllamaDefaultScheme, Registry: OllamaDefaultRegistry, Namespace: OllamaDefaultNamespace, Tag: OllamaDefaultTag, } + { + if o.DefaultScheme != "" { + om.Schema = o.DefaultScheme + } + if o.DefaultRegistry != "" { + om.Registry = o.DefaultRegistry + } + if o.DefaultNamespace != "" { + om.Namespace = o.DefaultNamespace + } + if o.DefaultTag != "" { + om.Tag = o.DefaultTag + } + } m := model diff --git a/ollama_model_option.go b/ollama_model_option.go new file mode 100644 index 0000000..325d4eb --- /dev/null +++ b/ollama_model_option.go @@ -0,0 +1,79 @@ +package gguf_parser + +import ( + "net/url" + "strings" +) + +type ( + _OllamaModelOptions struct { + DefaultScheme string + DefaultRegistry string + DefaultNamespace string + DefaultTag string + } + OllamaModelOption func(*_OllamaModelOptions) +) + +// SetOllamaModelBaseURL parses the given base URL, +// and sets default schema/registry for OllamaModel. +func SetOllamaModelBaseURL(baseURL string) OllamaModelOption { + baseURL = strings.TrimSpace(baseURL) + return func(o *_OllamaModelOptions) { + if baseURL == "" { + return + } + + if !strings.Contains(baseURL, "://") { + baseURL = "https://" + baseURL + } + + u, err := url.Parse(baseURL) + if err != nil { + return + } + + o.DefaultScheme = u.Scheme + o.DefaultRegistry = u.Host + } +} + +// SetOllamaModelDefaultScheme sets the default scheme for OllamaModel. +func SetOllamaModelDefaultScheme(scheme string) OllamaModelOption { + return func(o *_OllamaModelOptions) { + if scheme == "" { + return + } + o.DefaultScheme = scheme + } +} + +// SetOllamaModelDefaultRegistry sets the default registry for OllamaModel. +func SetOllamaModelDefaultRegistry(registry string) OllamaModelOption { + return func(o *_OllamaModelOptions) { + if registry == "" { + return + } + o.DefaultRegistry = registry + } +} + +// SetOllamaModelDefaultNamespace sets the default namespace for OllamaModel. +func SetOllamaModelDefaultNamespace(namespace string) OllamaModelOption { + return func(o *_OllamaModelOptions) { + if namespace == "" { + return + } + o.DefaultNamespace = namespace + } +} + +// SetOllamaModelDefaultTag sets the default tag for OllamaModel. +func SetOllamaModelDefaultTag(tag string) OllamaModelOption { + return func(o *_OllamaModelOptions) { + if tag == "" { + return + } + o.DefaultTag = tag + } +}