llmgw/internal/hf/hf.go

// Package hf reads the Hugging Face model registry.
//
// Enough of it to answer two questions an operator has before mirroring a
// model: what is there, and which file do I want. Nothing here downloads
// anything — the mirror does that, over its own network, and the point of the
// split is that the control plane never carries a model's bytes.
//
// Read-only and anonymous. A gated or private repository answers 401 to this
// client, which is reported rather than worked around: a token belongs to the
// machine that fetches the weights, not to the one browsing a catalogue.
package hf

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"regexp"
	"slices"
	"strings"
	"sync"
	"time"
)

// DefaultEndpoint is the public registry.
const DefaultEndpoint = "https://huggingface.co"

// Client reads the registry.
type Client struct {
	base string
	hc   *http.Client

	mu    sync.Mutex
	cache map[string]cached
}

// New builds a client against an endpoint, or the public one when empty.
func New(endpoint string, timeout time.Duration) *Client {
	if endpoint == "" {
		endpoint = DefaultEndpoint
	}
	if timeout <= 0 {
		// Short. This runs while somebody is looking at a page, and a registry
		// that is slow today is better reported than waited for.
		timeout = 10 * time.Second
	}
	return &Client{base: strings.TrimSuffix(endpoint, "/"),
		hc: &http.Client{Timeout: timeout}}
}

// Model is a repository as a search result shows it.
type Model struct {
	ID           string    `json:"id"`
	Downloads    int64     `json:"downloads"`
	Likes        int64     `json:"likes"`
	PipelineTag  string    `json:"pipeline_tag"`
	Tags         []string  `json:"tags"`
	LastModified time.Time `json:"lastModified"`
	// Gated is true when the repository asks for terms to be accepted. It
	// decodes from either a bool or a string ("auto", "manual"), which is why
	// it is not a plain bool here.
	Gated   Gated `json:"gated"`
	Private bool  `json:"private"`

	CreatedAt     time.Time `json:"createdAt"`
	TrendingScore float64   `json:"trendingScore"`
	// GGUF is what the registry read from the repository's weights, absent
	// for one that has none.
	GGUF *GGUFInfo `json:"gguf"`
	// Safetensors is the registry's summary of original weights.
	Safetensors *SafetensorsInfo `json:"safetensors,omitempty"`
	// Siblings are the repository's files, when asked for with WithFiles.
	Siblings []struct {
		RFilename string `json:"rfilename"`
	} `json:"siblings,omitempty"`
}

// SafetensorsInfo is the registry's summary of a repository's safetensors.
type SafetensorsInfo struct {
	Total int64 `json:"total"`
	// Parameters counts them by dtype: {"BF16": 8e9}. A packed quantisation
	// (AWQ, GPTQ, MLX) is counted as its I32 or U32 at the count of weights
	// it unpacks to, not of the integers that hold them.
	Parameters map[string]int64 `json:"parameters,omitempty"`
}

// dtypeBytes is what one parameter of a safetensors dtype takes.
var dtypeBytes = map[string]float64{
	"F64": 8, "I64": 8, "U64": 8, "F32": 4, "BF16": 2, "F16": 2, "I16": 2, "U16": 2,
	"F8_E4M3": 1, "F8_E5M2": 1, "F8_E8M0": 1, "I8": 1, "U8": 1, "BOOL": 1,
}

// WeightBytes is about how much the safetensors weights take: each dtype's
// count at its size, and a packed quantisation at 4 bits a weight, 8 for
// one tagged 8-bit — I32 or U32 always, U8 for MXFP4 and NVFP4, which the
// registry counts at the weights they unpack to. 0 when the registry has
// no summary.
func (m Model) WeightBytes() int64 {
	if m.Safetensors == nil {
		return 0
	}
	if len(m.Safetensors.Parameters) == 0 {
		return m.Safetensors.Total * 2
	}
	packed := 0.5
	if m.HasTag("8-bit") {
		packed = 1
	}
	fp4 := m.HasTag("mxfp4") || m.HasTag("nvfp4") || m.HasTag("fp4")
	var b float64
	for dt, n := range m.Safetensors.Parameters {
		switch per, ok := dtypeBytes[dt]; {
		case dt == "U8" && fp4:
			b += float64(n) * 0.5
		case ok:
			b += float64(n) * per
		case dt == "I32" || dt == "U32":
			b += float64(n) * packed
		default:
			b += float64(n) * 2
		}
	}
	return int64(b)
}

// FP8 reports whether the safetensors weights are in an 8-bit float, which
// vLLM runs only on cards of compute capability 8.9 or newer.
func (m Model) FP8() bool {
	if m.Safetensors == nil {
		return false
	}
	for dt := range m.Safetensors.Parameters {
		if strings.HasPrefix(dt, "F8_") {
			return true
		}
	}
	return false
}

// HasTag reports whether the registry tags the repository so.
func (m Model) HasTag(tag string) bool { return slices.Contains(m.Tags, tag) }

// GGUFInfo is the registry's summary of a repository's GGUF weights.
type GGUFInfo struct {
	// Total is the parameter count.
	Total         int64  `json:"total"`
	Architecture  string `json:"architecture"`
	ContextLength int64  `json:"context_length"`
}

// Params is the parameter count, 0 when the registry does not know it.
func (m Model) Params() int64 {
	switch {
	case m.GGUF != nil:
		return m.GGUF.Total
	case m.Safetensors != nil:
		return m.Safetensors.Total
	}
	return 0
}

// Architecture is the GGUF architecture, empty when unknown.
func (m Model) Architecture() string {
	if m.GGUF == nil {
		return ""
	}
	return m.GGUF.Architecture
}

// Gated is Hugging Face's gating flag, which is false or a word.
type Gated struct{ Reason string }

// UnmarshalJSON accepts false, true, or a string.
//
// The field is documented as a boolean and is a string in practice for most
// gated repositories. Decoding it as a bool fails the whole response, which
// would make one gated model in a page of results hide the other nineteen.
func (g *Gated) UnmarshalJSON(b []byte) error {
	var s string
	if err := json.Unmarshal(b, &s); err == nil {
		g.Reason = s
		return nil
	}
	var v bool
	if err := json.Unmarshal(b, &v); err != nil {
		return nil // unknown shape: treat as not gated rather than fail the page
	}
	if v {
		g.Reason = "yes"
	}
	return nil
}

// Is reports whether the repository is gated at all.
func (g Gated) Is() bool { return g.Reason != "" && g.Reason != "false" }

// Search finds repositories by name.
func (c *Client) Search(ctx context.Context, query string, limit int) ([]Model, error) {
	if strings.TrimSpace(query) == "" {
		return nil, fmt.Errorf("search for something")
	}
	// By downloads, because the thing somebody means by a short name is
	// almost always the one everybody else meant too.
	return c.Explore(ctx, Query{Search: query, Sort: SortDownloads, Limit: limit})
}

// Sort orders an exploration.
type Sort string

// The orders the registry offers that are worth offering.
const (
	SortTrending  Sort = "trendingScore"
	SortDownloads Sort = "downloads"
	SortLikes     Sort = "likes"
	SortCreated   Sort = "createdAt"
)

// Query is what to list from the registry.
type Query struct {
	Search string
	Sort   Sort
	// Task is a pipeline tag: text-generation, feature-extraction. Empty is
	// any.
	Task string
	// GGUF keeps to repositories with GGUF weights, which is what llama.cpp
	// loads.
	GGUF bool
	// MaxParams bounds the parameter count; 0 is no bound. The registry
	// applies it, so a page of small models is a full page.
	MaxParams int64
	Limit     int
	// Filters are further registry tags every result must carry, such as
	// "lora", or "base_model:adapter:Qwen/Qwen2.5-7B-Instruct" for the
	// adapters made for one model.
	Filters []string
	// WithFiles lists each repository's files too, to find one holding a
	// file by name.
	WithFiles bool
}

// exploreFields is everything a listing shows. Asking for any field with
// expand returns only the fields asked for, so all of them are named.
var exploreFields = []string{"downloads", "likes", "pipeline_tag", "tags", "lastModified",
	"createdAt", "gated", "private", "trendingScore", "gguf", "safetensors"}

// Explore lists repositories: the trending ones, the newest, or a search.
//
// Remembered for a few minutes. The mirror page opens on this, and a page
// that waits on a registry on the far side of the internet every time it is
// opened is slow for news that changes by the hour.
func (c *Client) Explore(ctx context.Context, q Query) ([]Model, error) {
	if q.Limit <= 0 || q.Limit > 100 {
		q.Limit = 20
	}
	if q.Sort == "" {
		q.Sort = SortTrending
	}
	v := url.Values{
		"limit":     {fmt.Sprint(q.Limit)},
		"sort":      {string(q.Sort)},
		"direction": {"-1"},
		"expand[]":  exploreFields,
	}
	if q.WithFiles {
		v.Add("expand[]", "siblings")
	}
	if s := strings.TrimSpace(q.Search); s != "" {
		v.Set("search", s)
	}
	if q.Task != "" {
		v.Set("pipeline_tag", q.Task)
	}
	if q.GGUF {
		v.Add("filter", "gguf")
	}
	for _, f := range q.Filters {
		v.Add("filter", f)
	}
	if q.MaxParams > 0 {
		v.Set("num_parameters", fmt.Sprintf("min:0,max:%d", q.MaxParams))
	}
	path := "/api/models?" + v.Encode()

	c.mu.Lock()
	if e, ok := c.cache[path]; ok && time.Since(e.at) < exploreTTL {
		c.mu.Unlock()
		return e.models, nil
	}
	c.mu.Unlock()

	var out []Model
	if err := c.get(ctx, path, &out); err != nil {
		return nil, err
	}
	c.mu.Lock()
	if c.cache == nil || len(c.cache) > 200 {
		c.cache = map[string]cached{}
	}
	c.cache[path] = cached{at: time.Now(), models: out}
	c.mu.Unlock()
	return out, nil
}

// exploreTTL is how long a listing is reused.
const exploreTTL = 10 * time.Minute

type cached struct {
	at     time.Time
	models []Model
}

// File is one file in a repository.
type File struct {
	Path string
	// Bytes is the real size. For a file stored in LFS — every weights file —
	// the plain size field is the size of the pointer, a few hundred bytes, so
	// reading that one would report a 5 GB model as 133 bytes.
	Bytes int64
	// SHA256 is what the registry publishes for an LFS file, so it can be
	// compared with what the mirror computes after fetching. Empty for a small
	// file stored in git, where the oid is a git blob hash and not a checksum
	// of the content.
	SHA256 string
}

// IsWeights reports whether this file is a model rather than documentation.
//
// By extension, which is crude and right: the registry does not label them, and
// an operator filling a mirror wants the four files out of thirty that are
// worth gigabytes.
func (f File) IsWeights() bool {
	for _, ext := range []string{".gguf", ".safetensors", ".bin", ".pt", ".pth"} {
		if strings.HasSuffix(strings.ToLower(f.Path), ext) {
			return true
		}
	}
	return false
}

// companionNames are the files beside safetensors weights that an engine
// loading the directory reads: the model's shape, its tokenizer, how it was
// quantised, and which shard holds which tensor.
var companionNames = map[string]bool{
	"config.json": true, "generation_config.json": true,
	"tokenizer.json": true, "tokenizer_config.json": true, "tokenizer.model": true,
	"special_tokens_map.json": true, "added_tokens.json": true,
	"vocab.json": true, "vocab.txt": true, "merges.txt": true,
	"chat_template.jinja": true, "chat_template.json": true,
	"preprocessor_config.json": true, "processor_config.json": true,
	"quantize_config.json": true, "quant_config.json": true,
}

// Companions are the files of a repository an engine reads beside a
// safetensors set: config.json and the tokenizer, the shard index. Only
// those at the top: the directory an engine is given is one level deep.
//
// Never code (a .py): a model that needs its own Python to load asks the
// engine to run what the repository ships, which is a decision for whoever
// runs the engine, made on purpose, not a file swept in with the tokenizer.
func Companions(files []File) []File {
	var out []File
	for _, f := range files {
		name := f.Path
		if strings.Contains(name, "/") {
			continue
		}
		if companionNames[name] || strings.HasSuffix(name, ".safetensors.index.json") ||
			strings.HasSuffix(name, ".tiktoken") {
			out = append(out, f)
		}
	}
	return out
}

// Repo is a repository's current revision and what is in it.
type Repo struct {
	ID string
	// Revision is the commit the listing came from.
	//
	// Carried so a mirrored copy can name it. Fetching from "main" records
	// nothing about what was fetched, and a repository that is re-uploaded
	// under the same branch leaves two different sets of weights with the same
	// provenance and no way to tell them apart.
	Revision string
	Gated    Gated
	Files    []File
	// Architecture is the GGUF architecture the registry read from the
	// repository's weights, empty for a repository with none. What decides
	// whether llama.cpp can load it, known before anything is downloaded.
	Architecture string
	// Params is the parameter count the registry read from the weights, 0
	// when it did not.
	Params int64

	// From the model card's header: what somebody asks before reading the
	// rest of it.
	License   string
	BaseModel []string
	Languages []string
	// Card is the model card's text, markdown, without its header. Empty when
	// the repository has none or it could not be read — a repository is
	// still worth fetching from without its README.
	Card string
}

// ShortRevision is the revision as a person would write it.
func (r Repo) ShortRevision() string {
	if len(r.Revision) > 8 {
		return r.Revision[:8]
	}
	return r.Revision
}

// Base is the registry's address, for building a file's download URL.
func (c *Client) Base() string { return c.base }

// Repo reads a repository's revision and file list.
func (c *Client) Repo(ctx context.Context, id string) (Repo, error) { return c.RepoAt(ctx, id, "") }

// RepoAt reads a repository as it was at a commit, which may be shortened;
// "" is its latest.
func (c *Client) RepoAt(ctx context.Context, id, rev string) (Repo, error) {
	id = strings.Trim(strings.TrimSpace(id), "/")
	if !strings.Contains(id, "/") {
		return Repo{}, fmt.Errorf("%q is not a repository; it looks like owner/name", id)
	}

	var meta struct {
		ID    string `json:"id"`
		SHA   string `json:"sha"`
		Gated Gated  `json:"gated"`
		GGUF  *struct {
			Architecture string `json:"architecture"`
			Total        int64  `json:"total"`
		} `json:"gguf"`
		Safetensors *struct {
			Total int64 `json:"total"`
		} `json:"safetensors"`
		CardData struct {
			License   string    `json:"license"`
			BaseModel stringish `json:"base_model"`
			Language  stringish `json:"language"`
		} `json:"cardData"`
	}
	// Expanded to what this reads: asking for gguf alone would leave out the
	// revision and the gate.
	at := ""
	if rev != "" {
		at = "/revision/" + url.PathEscape(rev)
	}
	if err := c.get(ctx, "/api/models/"+id+at+"?expand[]=sha&expand[]=gated&expand[]=gguf&expand[]=safetensors&expand[]=cardData", &meta); err != nil {
		return Repo{}, err
	}

	var tree []struct {
		Type string `json:"type"`
		Path string `json:"path"`
		Size int64  `json:"size"`
		LFS  *struct {
			OID  string `json:"oid"`
			Size int64  `json:"size"`
		} `json:"lfs"`
	}
	// By revision rather than by branch, so the listing and the URLs built from
	// it describe the same commit even if the branch moves in between.
	rev = meta.SHA
	if rev == "" {
		rev = "main"
	}
	// Recursive: a repository that keeps each quantisation in a directory of
	// its own — ISTA-DASLab's IQ2_XS/, IQ3_S/ — listed no weights at all.
	if err := c.get(ctx, "/api/models/"+id+"/tree/"+rev+"?recursive=true", &tree); err != nil {
		return Repo{}, err
	}

	r := Repo{ID: meta.ID, Revision: meta.SHA, Gated: meta.Gated,
		License: meta.CardData.License, BaseModel: meta.CardData.BaseModel,
		Languages: meta.CardData.Language}
	if meta.GGUF != nil {
		r.Architecture, r.Params = meta.GGUF.Architecture, meta.GGUF.Total
	}
	if meta.Safetensors != nil && meta.Safetensors.Total > 0 {
		r.Params = meta.Safetensors.Total
	}
	if r.ID == "" {
		r.ID = id
	}
	hasCard := false
	for _, e := range tree {
		if e.Type != "file" {
			continue
		}
		if e.Path == "README.md" {
			hasCard = true
		}
		f := File{Path: e.Path, Bytes: e.Size}
		if e.LFS != nil {
			f.Bytes, f.SHA256 = e.LFS.Size, e.LFS.OID
		}
		r.Files = append(r.Files, f)
	}
	if hasCard {
		// Best effort: the files are what the page is for.
		if b, err := c.raw(ctx, "/"+r.ID+"/raw/"+rev+"/README.md", maxCard); err == nil {
			r.Card = stripFrontMatter(string(b))
		}
	}
	return r, nil
}

// maxCard is as much of a model card as is read. Most are a few kilobytes;
// the cap is for the one that is not.
const maxCard = 256 << 10

// stripFrontMatter drops a card's YAML header, which the registry already
// hands over as cardData and which reads as noise at the top of the text.
func stripFrontMatter(s string) string {
	s = strings.TrimPrefix(s, "\ufeff")
	if !strings.HasPrefix(s, "---") {
		return strings.TrimSpace(s)
	}
	rest := s[3:]
	if i := strings.Index(rest, "\n---"); i >= 0 {
		rest = rest[i+4:]
		if j := strings.IndexByte(rest, '\n'); j >= 0 {
			rest = rest[j+1:]
		} else {
			rest = ""
		}
		return strings.TrimSpace(rest)
	}
	return strings.TrimSpace(s)
}

// stringish is a card field written as one string or a list of them.
type stringish []string

// UnmarshalJSON accepts a string, a list of strings, or anything else as
// nothing: a card header is free-form YAML, and one odd field should not
// cost the page its file list.
func (v *stringish) UnmarshalJSON(b []byte) error {
	var one string
	if json.Unmarshal(b, &one) == nil {
		if one != "" {
			*v = []string{one}
		}
		return nil
	}
	var many []string
	if json.Unmarshal(b, &many) == nil {
		*v = many
	}
	return nil
}

// BaseModel is the model a repository was made from, read from its tags:
// "base_model:quantized:Qwen/Qwen3-8B" and its kin. Empty when not said.
func (m Model) BaseModel() string {
	for _, t := range m.Tags {
		if !strings.HasPrefix(t, "base_model:") {
			continue
		}
		rest := strings.TrimPrefix(t, "base_model:")
		// The relation is optional: base_model:X or base_model:quantized:X.
		if i := strings.Index(rest, ":"); i >= 0 && !strings.Contains(rest[:i], "/") {
			rest = rest[i+1:]
		}
		if rest != "" {
			return rest
		}
	}
	return ""
}

// License is the license tag, empty when there is none.
func (m Model) License() string {
	for _, t := range m.Tags {
		if strings.HasPrefix(t, "license:") {
			return strings.TrimPrefix(t, "license:")
		}
	}
	return ""
}

// DownloadURL is where the mirror should fetch a file from.
//
// Pinned to the revision the listing came from. "main" would mean the mirror
// fetches whatever is there when it gets round to it, which is not necessarily
// what was on screen when somebody pressed the button.
func (r Repo) DownloadURL(base, path string) string {
	if base == "" {
		base = DefaultEndpoint
	}
	rev := r.Revision
	if rev == "" {
		rev = "main"
	}
	return strings.TrimSuffix(base, "/") + "/" + r.ID + "/resolve/" + rev + "/" + path
}

func (c *Client) get(ctx context.Context, path string, out any) error {
	body, err := c.raw(ctx, path, 4<<20)
	if err != nil {
		return err
	}
	if err := json.Unmarshal(body, out); err != nil {
		return fmt.Errorf("the registry's answer was not the JSON expected: %w", err)
	}
	return nil
}

// raw reads a path from the registry, at most limit bytes of it.
func (c *Client) raw(ctx context.Context, path string, limit int64) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.base+path, nil)
	if err != nil {
		return nil, err
	}
	resp, err := c.hc.Do(req)
	if err != nil {
		return nil, fmt.Errorf("reaching the model registry: %w", err)
	}
	defer resp.Body.Close()

	// Capped: this is a third party, and a page that renders what it sends
	// should not be handed a gigabyte of it.
	body, err := io.ReadAll(io.LimitReader(resp.Body, limit))
	if err != nil {
		return nil, err
	}
	switch {
	case resp.StatusCode == http.StatusNotFound:
		return nil, fmt.Errorf("the registry has nothing at %s", path)
	case resp.StatusCode == http.StatusUnauthorized, resp.StatusCode == http.StatusForbidden:
		// Said plainly, because the next thing somebody does is check their
		// network. It is the repository, not the connection.
		return nil, fmt.Errorf("the registry refused this anonymously; the repository is gated or private")
	case resp.StatusCode >= 300:
		return nil, fmt.Errorf("the registry answered %d", resp.StatusCode)
	}
	return body, nil
}

// shardPattern matches the "-00002-of-00004" a split file carries before its
// extension. Five digits in practice, but not required: the count is what
// matters and a repository that used three would otherwise list every part as
// a separate model.
var shardPattern = regexp.MustCompile(`^(.*)-(\d+)-of-(\d+)(\.[^.]+)$`)

// FileSet is one thing worth fetching: a file, or every shard of one.
//
// The distinction the page turns on. Two GGUF files in a repository are usually
// alternatives — q4_k_m or q8_0, you want one — while two safetensors files are
// halves of the same weights and useless apart. Fetching "all the weights" is
// right for the second and wasteful for the first, so neither is offered:
// a set is offered, and a set is whole.
type FileSet struct {
	// Name is what to call it. For shards it keeps the -of- notation, because
	// a repository can hold both the split and the joined copy of the same
	// quantisation and they must not read as one entry.
	Name   string
	Files  []File
	Bytes  int64
	Shards int
}

// SHA256 is the checksum when there is one to give.
//
// Empty for a shard set: several files have several checksums and no single
// value for the set, and showing one of them invites comparing it against the
// wrong thing.
func (s FileSet) SHA256() string {
	if len(s.Files) == 1 {
		return s.Files[0].SHA256
	}
	return ""
}

// Group collects shards into sets, leaving every other file on its own.
//
// Order is preserved from the listing, so a repository that sorts its files
// sensibly still reads sensibly here.
func Group(files []File) []FileSet {
	var out []FileSet
	at := map[string]int{}
	for _, f := range files {
		m := shardPattern.FindStringSubmatch(f.Path)
		if m == nil {
			out = append(out, FileSet{Name: f.Path, Files: []File{f}, Bytes: f.Bytes})
			continue
		}
		prefix, total, ext := m[1], m[3], m[4]
		name := prefix + "-*-of-" + total + ext
		i, ok := at[name]
		if !ok {
			at[name] = len(out)
			out = append(out, FileSet{Name: name})
			i = at[name]
		}
		out[i].Files = append(out[i].Files, f)
		out[i].Bytes += f.Bytes
		out[i].Shards = len(out[i].Files)
	}
	return out
}

// ModelName is the catalog name a repository suggests: its last element,
// lowercased, with the -GGUF marker dropped. The owner is not part of what
// anybody calls the model, and "-GGUF" describes the file format.
func ModelName(repoID string) string {
	name := repoID
	if i := strings.LastIndex(name, "/"); i >= 0 {
		name = name[i+1:]
	}
	name = strings.TrimSuffix(name, "-GGUF")
	name = strings.TrimSuffix(name, "-gguf")
	return strings.ToLower(name)
}

// FileSource reads a source naming one file of a registry repository: hf:
// with the file, or a resolve URL on the registry. A resolve URL already at
// a commit is still looked up, for the file's sha256.
func FileSource(src, base string) (id, rev, file string, ok bool) {
	if rest, found := strings.CutPrefix(src, "hf:"); found {
		parts := strings.Split(rest, "/")
		if len(parts) < 3 {
			return "", "", "", false
		}
		return parts[0] + "/" + parts[1], "", strings.Join(parts[2:], "/"), true
	}
	u, err := url.Parse(src)
	if err != nil || (u.Scheme != "https" && u.Scheme != "http") || !isRegistryHost(u.Host, base) {
		return "", "", "", false
	}
	parts := strings.Split(strings.Trim(u.Path, "/"), "/")
	if len(parts) < 5 || parts[2] != "resolve" {
		return "", "", "", false
	}
	return parts[0] + "/" + parts[1], parts[3], strings.Join(parts[4:], "/"), true
}

// isRegistryHost says whether a host is the registry's: Hugging Face's own, or
// the endpoint the cluster names instead.
func isRegistryHost(host, base string) bool {
	if host == "huggingface.co" {
		return true
	}
	b, err := url.Parse(base)
	return err == nil && b.Host != "" && b.Host == host
}