chore: import upstream snapshot with attribution

internal/embedding/potion.go

@@ -0,0 +1,321 @@
+package embedding
+
+import (
+	"context"
+	"crypto/sha256"
+	"encoding/binary"
+	"encoding/hex"
+	"encoding/json"
+	"fmt"
+	"io"
+	"math"
+	"net/http"
+	"os"
+	"path/filepath"
+	"strings"
+	"time"
+
+	"github.com/zzet/gortex/internal/platform"
+)
+
+// PotionProvider embeds text with a Model2Vec static embedding model —
+// a WordPiece tokenizer plus a single [vocab × dims] embedding matrix.
+// Inference is: tokenize → gather the token rows → mean-pool →
+// L2-normalise (the model config's `normalize: true` convention, matching
+// the reference implementations). No transformer, no positional limit —
+// a 50-candidate rerank batch embeds in well under a millisecond.
+//
+// The bundled model is minishlab/potion-code-16M-v2 (MIT): 256-dim
+// vectors over a ~63.5k-token code-mined vocabulary, distilled from a
+// code-retrieval teacher. It replaces the older averaged word-vector
+// channel for the rerank's semantic-cosine signal.
+type PotionProvider struct {
+	tok  *wordPieceTokenizer
+	mat  []float32 // row-major [vocab][dims]
+	dims int
+}
+
+// Potion model pin. The revision and checksums identify the exact
+// artifacts the loader accepts; a mismatched download is discarded.
+const (
+	potionModelName = "potion-code-16M-v2"
+	potionRevision  = "d3daf3e31f36d78f75913030b8bdf4a505d5b833"
+	potionBaseURL   = "https://huggingface.co/minishlab/" + potionModelName + "/resolve/" + potionRevision
+
+	potionWeightsFile   = "model.safetensors"
+	potionTokenizerFile = "tokenizer.json"
+
+	potionWeightsSHA256   = "75cf7a6c2171b230ad19b1e7d8e0b1aee86da5a02af8e7cacedd9921d227623c"
+	potionTokenizerSHA256 = "107bbdcbad4bff1d299b7a4c3a2fb17c52890688b7dd0e4c9deab79d3c4f3d45"
+)
+
+// maxPotionTokens caps how many tokens of one text feed the mean-pool.
+// The model is static (no sequence limit), but the rerank only ever
+// embeds short name+signature+doc fragments; the cap bounds the cost of
+// a pathological input.
+const maxPotionTokens = 512
+
+// NewPotionProviderFromDir loads the model from a directory holding
+// model.safetensors and tokenizer.json.
+func NewPotionProviderFromDir(dir string) (*PotionProvider, error) {
+	tok, err := loadWordPieceTokenizer(filepath.Join(dir, potionTokenizerFile))
+	if err != nil {
+		return nil, err
+	}
+	mat, dims, err := loadSafetensorsF16(filepath.Join(dir, potionWeightsFile))
+	if err != nil {
+		return nil, err
+	}
+	return &PotionProvider{tok: tok, mat: mat, dims: dims}, nil
+}
+
+func (p *PotionProvider) Dimensions() int { return p.dims }
+func (p *PotionProvider) Close() error    { return nil }
+
+func (p *PotionProvider) Embed(_ context.Context, text string) ([]float32, error) {
+	return p.embed(text), nil
+}
+
+func (p *PotionProvider) EmbedBatch(_ context.Context, texts []string) ([][]float32, error) {
+	out := make([][]float32, len(texts))
+	for i, t := range texts {
+		out[i] = p.embed(t)
+	}
+	return out, nil
+}
+
+func (p *PotionProvider) embed(text string) []float32 {
+	ids := p.tok.Encode(text)
+	if len(ids) > maxPotionTokens {
+		ids = ids[:maxPotionTokens]
+	}
+	vec := make([]float32, p.dims)
+	if len(ids) == 0 {
+		return vec
+	}
+	rows := len(p.mat) / p.dims
+	n := 0
+	for _, id := range ids {
+		if id < 0 || id >= rows {
+			continue
+		}
+		row := p.mat[id*p.dims : (id+1)*p.dims]
+		for i, v := range row {
+			vec[i] += v
+		}
+		n++
+	}
+	if n == 0 {
+		return vec
+	}
+	inv := 1.0 / float32(n)
+	var norm float64
+	for i := range vec {
+		vec[i] *= inv
+		norm += float64(vec[i]) * float64(vec[i])
+	}
+	// L2-normalise (model config `normalize: true`).
+	if norm > 0 {
+		s := float32(1.0 / math.Sqrt(norm))
+		for i := range vec {
+			vec[i] *= s
+		}
+	}
+	return vec
+}
+
+// loadSafetensorsF16 reads a single-tensor safetensors file holding an
+// F16 "embeddings" matrix and returns it as row-major float32.
+func loadSafetensorsF16(path string) ([]float32, int, error) {
+	raw, err := os.ReadFile(path)
+	if err != nil {
+		return nil, 0, fmt.Errorf("read weights: %w", err)
+	}
+	if len(raw) < 8 {
+		return nil, 0, fmt.Errorf("weights file too short")
+	}
+	hlen := binary.LittleEndian.Uint64(raw[:8])
+	if hlen == 0 || 8+hlen > uint64(len(raw)) {
+		return nil, 0, fmt.Errorf("weights header out of range")
+	}
+	var header map[string]json.RawMessage
+	if err := json.Unmarshal(raw[8:8+hlen], &header); err != nil {
+		return nil, 0, fmt.Errorf("parse weights header: %w", err)
+	}
+	var info struct {
+		Dtype       string `json:"dtype"`
+		Shape       []int  `json:"shape"`
+		DataOffsets []int  `json:"data_offsets"`
+	}
+	entry, ok := header["embeddings"]
+	if !ok {
+		return nil, 0, fmt.Errorf("weights file has no 'embeddings' tensor")
+	}
+	if err := json.Unmarshal(entry, &info); err != nil {
+		return nil, 0, fmt.Errorf("parse tensor info: %w", err)
+	}
+	if info.Dtype != "F16" {
+		return nil, 0, fmt.Errorf("unsupported tensor dtype %q (want F16)", info.Dtype)
+	}
+	if len(info.Shape) != 2 || len(info.DataOffsets) != 2 {
+		return nil, 0, fmt.Errorf("unexpected tensor shape/offsets")
+	}
+	rows, dims := info.Shape[0], info.Shape[1]
+	start := 8 + int(hlen) + info.DataOffsets[0]
+	end := 8 + int(hlen) + info.DataOffsets[1]
+	if start < 0 || end > len(raw) || end-start != rows*dims*2 {
+		return nil, 0, fmt.Errorf("tensor data out of range")
+	}
+	data := raw[start:end]
+	mat := make([]float32, rows*dims)
+	for i := range mat {
+		mat[i] = f16ToF32(binary.LittleEndian.Uint16(data[i*2 : i*2+2]))
+	}
+	return mat, dims, nil
+}
+
+// f16ToF32 converts an IEEE-754 half-precision value to float32.
+func f16ToF32(h uint16) float32 {
+	sign := uint32(h>>15) << 31
+	exp := uint32(h>>10) & 0x1f
+	man := uint32(h) & 0x3ff
+	switch exp {
+	case 0:
+		if man == 0 {
+			return math.Float32frombits(sign)
+		}
+		// Subnormal half: renormalise into a normal float32.
+		e := uint32(113) // 127 - 14
+		for man&0x400 == 0 {
+			man <<= 1
+			e--
+		}
+		man &= 0x3ff
+		return math.Float32frombits(sign | e<<23 | man<<13)
+	case 0x1f:
+		return math.Float32frombits(sign | 0xff<<23 | man<<13)
+	default:
+		return math.Float32frombits(sign | (exp+112)<<23 | man<<13)
+	}
+}
+
+// resolvePotionDir returns the first directory that holds both model
+// files, probing in order:
+//
+//  1. $GORTEX_POTION_DIR — explicit override.
+//  2. <executable dir>/models/<model> — the release/bench sidecar. A
+//     packaged install ships the model next to the binary, so it is
+//     fully offline after install.
+//  3. <gortex home>/models/<model> — the per-user cache the first-use
+//     download fills.
+//
+// Returns "" when none has the files.
+func resolvePotionDir() string {
+	var candidates []string
+	if d := os.Getenv("GORTEX_POTION_DIR"); d != "" {
+		candidates = append(candidates, d)
+	}
+	if exe, err := os.Executable(); err == nil {
+		candidates = append(candidates, filepath.Join(filepath.Dir(exe), "models", potionModelName))
+	}
+	candidates = append(candidates, filepath.Join(platform.ModelsDir(), potionModelName))
+	for _, dir := range candidates {
+		if hasPotionFiles(dir) {
+			return dir
+		}
+	}
+	return ""
+}
+
+func hasPotionFiles(dir string) bool {
+	for _, f := range []string{potionWeightsFile, potionTokenizerFile} {
+		if st, err := os.Stat(filepath.Join(dir, f)); err != nil || st.IsDir() {
+			return false
+		}
+	}
+	return true
+}
+
+// downloadPotion fetches the pinned model revision into the per-user
+// models dir, verifying each file's SHA-256 before moving it into
+// place. Returns the directory on success. Disabled entirely when
+// GORTEX_POTION_DOWNLOAD=0 (offline installs ship the sidecar instead).
+func downloadPotion() (string, error) {
+	if v := strings.TrimSpace(os.Getenv("GORTEX_POTION_DOWNLOAD")); v == "0" || strings.EqualFold(v, "false") || strings.EqualFold(v, "off") {
+		return "", fmt.Errorf("model download disabled by GORTEX_POTION_DOWNLOAD")
+	}
+	// Test binaries never reach the network: a `go test` run that
+	// exercises a search path must stay hermetic and fall back to the
+	// baked vectors instead of pulling 32MB mid-suite.
+	if strings.HasSuffix(os.Args[0], ".test") {
+		return "", fmt.Errorf("model download disabled in test binaries")
+	}
+	dir := filepath.Join(platform.ModelsDir(), potionModelName)
+	if err := os.MkdirAll(dir, 0o755); err != nil {
+		return "", err
+	}
+	client := &http.Client{Timeout: 120 * time.Second}
+	files := []struct{ name, sum string }{
+		{potionWeightsFile, potionWeightsSHA256},
+		{potionTokenizerFile, potionTokenizerSHA256},
+	}
+	for _, f := range files {
+		dst := filepath.Join(dir, f.name)
+		if st, err := os.Stat(dst); err == nil && !st.IsDir() {
+			if ok, _ := fileSHA256Matches(dst, f.sum); ok {
+				continue
+			}
+		}
+		if err := downloadVerified(client, potionBaseURL+"/"+f.name, dst, f.sum); err != nil {
+			return "", fmt.Errorf("fetch %s: %w", f.name, err)
+		}
+	}
+	return dir, nil
+}
+
+// downloadVerified streams url into dst.partial, verifies the SHA-256,
+// and renames into place atomically.
+func downloadVerified(client *http.Client, url, dst, wantSHA string) error {
+	resp, err := client.Get(url)
+	if err != nil {
+		return err
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("HTTP %d", resp.StatusCode)
+	}
+	tmp := dst + ".partial"
+	out, err := os.Create(tmp)
+	if err != nil {
+		return err
+	}
+	h := sha256.New()
+	_, cpErr := io.Copy(io.MultiWriter(out, h), resp.Body)
+	closeErr := out.Close()
+	if cpErr != nil {
+		os.Remove(tmp)
+		return cpErr
+	}
+	if closeErr != nil {
+		os.Remove(tmp)
+		return closeErr
+	}
+	if got := hex.EncodeToString(h.Sum(nil)); got != wantSHA {
+		os.Remove(tmp)
+		return fmt.Errorf("checksum mismatch: got %s want %s", got, wantSHA)
+	}
+	return os.Rename(tmp, dst)
+}
+
+func fileSHA256Matches(path, want string) (bool, error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return false, err
+	}
+	defer f.Close()
+	h := sha256.New()
+	if _, err := io.Copy(h, f); err != nil {
+		return false, err
+	}
+	return hex.EncodeToString(h.Sum(nil)) == want, nil
+}