feat: support HF_ENDPOINT/MS_ENDPOINT

Signed-off-by: thxCode <[email protected]>
gpustack · Nov 19, 2024 · bb8aaa2 · bb8aaa2
1 parent 0f0a056
commit bb8aaa2
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 23 deletions.
diff --git a/README.md b/README.md
@@ -237,6 +237,9 @@ $ gguf-parser --url="https://huggingface.co/MaziyarPanahi/Meta-Llama-3.1-405B-In
 
 #### Parse From HuggingFace
 
+> ![NOTE]
+> Allow using `HF_ENDPOINT` to override the default HuggingFace endpoint: `https://huggingface.co`.
+
 ```shell
 $ gguf-parser --hf-repo="openbmb/MiniCPM-Llama3-V-2_5-gguf" --hf-file="ggml-model-Q5_K_M.gguf" --hf-mmproj-file="mmproj-model-f16.gguf"
 +-------------------------------------------------------------------------------------------+
@@ -314,6 +317,9 @@ $ gguf-parser --hf-repo="etemiz/Llama-3.1-405B-Inst-GGUF" --hf-file="llama-3.1-4
 
 #### Parse From ModelScope
 
+> ![NOTE]
+> Allow using `MS_ENDPOINT` to override the default ModelScope endpoint: `https://modelscope.cn`.
+
 ```shell
 $ gguf-parser --ms-repo="shaowenchen/chinese-alpaca-2-13b-16k-gguf" --ms-file="chinese-alpaca-2-13b-16k.Q5_K.gguf"
 +------------------------------------------------------------------------------------------+
@@ -354,6 +360,9 @@ $ gguf-parser --ms-repo="shaowenchen/chinese-alpaca-2-13b-16k-gguf" --ms-file="c
 
 #### Parse From Ollama Library
 
+> ![NOTE]
+> Allow using `--ol-base-url` to override the default Ollama registry endpoint: `https://registry.ollama.ai`.
+
 ```shell
 $ gguf-parser --ol-model="llama3.1"
 +-----------------------------------------------------------------------------------------------------------+

diff --git a/file_from_remote.go b/file_from_remote.go
@@ -6,6 +6,7 @@ import (
 	"io"
 	"net/http"
 	"path/filepath"
+	"strings"
 	"time"
 
 	"github.com/gpustack/gguf-parser-go/util/httpx"
@@ -15,14 +16,16 @@ import (
 // ParseGGUFFileFromHuggingFace parses a GGUF file from Hugging Face(https://huggingface.co/),
 // and returns a GGUFFile, or an error if any.
 func ParseGGUFFileFromHuggingFace(ctx context.Context, repo, file string, opts ...GGUFReadOption) (*GGUFFile, error) {
-	return ParseGGUFFileRemote(ctx, fmt.Sprintf("https://huggingface.co/%s/resolve/main/%s", repo, file), opts...)
+	ep := osx.Getenv("HF_ENDPOINT", "https://huggingface.co")
+	return ParseGGUFFileRemote(ctx, fmt.Sprintf("%s/%s/resolve/main/%s", ep, repo, file), opts...)
 }
 
 // ParseGGUFFileFromModelScope parses a GGUF file from Model Scope(https://modelscope.cn/),
 // and returns a GGUFFile, or an error if any.
 func ParseGGUFFileFromModelScope(ctx context.Context, repo, file string, opts ...GGUFReadOption) (*GGUFFile, error) {
+	ep := osx.Getenv("MS_ENDPOINT", "https://modelscope.cn")
 	opts = append(opts[:len(opts):len(opts)], SkipRangeDownloadDetection())
-	return ParseGGUFFileRemote(ctx, fmt.Sprintf("https://modelscope.cn/models/%s/resolve/master/%s", repo, file), opts...)
+	return ParseGGUFFileRemote(ctx, fmt.Sprintf("%s/models/%s/resolve/master/%s", ep, repo, file), opts...)
 }
 
 // ParseGGUFFileRemote parses a GGUF file from a remote BlobURL,
@@ -59,31 +62,45 @@ func ParseGGUFFileRemote(ctx context.Context, url string, opts ...GGUFReadOption
 	cli := httpx.Client(
 		httpx.ClientOptions().
 			WithUserAgent("gguf-parser-go").
-			If(o.Debug, func(x *httpx.ClientOption) *httpx.ClientOption {
-				return x.WithDebug()
-			}).
-			If(o.BearerAuthToken != "", func(x *httpx.ClientOption) *httpx.ClientOption {
-				return x.WithBearerAuth(o.BearerAuthToken)
-			}).
+			If(o.Debug,
+				func(x *httpx.ClientOption) *httpx.ClientOption {
+					return x.WithDebug()
+				},
+			).
+			If(o.BearerAuthToken != "",
+				func(x *httpx.ClientOption) *httpx.ClientOption {
+					return x.WithBearerAuth(o.BearerAuthToken)
+				},
+			).
 			WithTimeout(0).
 			WithTransport(
 				httpx.TransportOptions().
 					WithoutKeepalive().
 					TimeoutForDial(5*time.Second).
 					TimeoutForTLSHandshake(5*time.Second).
 					TimeoutForResponseHeader(5*time.Second).
-					If(o.SkipProxy, func(x *httpx.TransportOption) *httpx.TransportOption {
-						return x.WithoutProxy()
-					}).
-					If(o.ProxyURL != nil, func(x *httpx.TransportOption) *httpx.TransportOption {
-						return x.WithProxy(http.ProxyURL(o.ProxyURL))
-					}).
-					If(o.SkipTLSVerification, func(x *httpx.TransportOption) *httpx.TransportOption {
-						return x.WithoutInsecureVerify()
-					}).
-					If(o.SkipDNSCache, func(x *httpx.TransportOption) *httpx.TransportOption {
-						return x.WithoutDNSCache()
-					})))
+					If(o.SkipProxy,
+						func(x *httpx.TransportOption) *httpx.TransportOption {
+							return x.WithoutProxy()
+						},
+					).
+					If(o.ProxyURL != nil,
+						func(x *httpx.TransportOption) *httpx.TransportOption {
+							return x.WithProxy(http.ProxyURL(o.ProxyURL))
+						},
+					).
+					If(o.SkipTLSVerification || !strings.HasPrefix(url, "https://"),
+						func(x *httpx.TransportOption) *httpx.TransportOption {
+							return x.WithoutInsecureVerify()
+						},
+					).
+					If(o.SkipDNSCache,
+						func(x *httpx.TransportOption) *httpx.TransportOption {
+							return x.WithoutDNSCache()
+						},
+					),
+			),
+	)
 
 	return parseGGUFFileFromRemote(ctx, cli, url, o)
 }
@@ -115,9 +132,12 @@ func parseGGUFFileFromRemote(ctx context.Context, cli *http.Client, url string,
 		sf, err := httpx.OpenSeekerFile(cli, req,
 			httpx.SeekerFileOptions().
 				WithBufferSize(o.BufferSize).
-				If(o.SkipRangeDownloadDetection, func(x *httpx.SeekerFileOption) *httpx.SeekerFileOption {
-					return x.WithoutRangeDownloadDetect()
-				}))
+				If(o.SkipRangeDownloadDetection,
+					func(x *httpx.SeekerFileOption) *httpx.SeekerFileOption {
+						return x.WithoutRangeDownloadDetect()
+					},
+				),
+		)
 		if err != nil {
 			return nil, fmt.Errorf("open http file: %w", err)
 		}