claude
f8ec77de8a
Go tool that computes e5-small embeddings for every corpus row, assigns frozen/dev split and grouped CV folds, outputs JSON.
274 lines
8.4 KiB
Go
274 lines
8.4 KiB
Go
// semantic-router-experiment computes embeddings for the semantic coarse-route
|
|
// corpus using the deployed multilingual-e5-small ONNX model. It outputs a
|
|
// JSON file containing every corpus row with its embedding vector, split
|
|
// assignment, and fold membership for grouped cross-validation.
|
|
//
|
|
// Usage:
|
|
//
|
|
// MAVEN_ONNX_LIB=/path/to/libonnxruntime.so \
|
|
// go run ./cmd/semantic-router-experiment/ -out embeddings.json
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"flag"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"sort"
|
|
"strings"
|
|
|
|
"github.com/kami/maven/internal/router"
|
|
"github.com/kami/maven/internal/router/semantic"
|
|
)
|
|
|
|
// CachedRow is one corpus row with its precomputed embedding and split metadata.
|
|
type CachedRow struct {
|
|
Text string `json:"text"`
|
|
Route string `json:"route"`
|
|
Source string `json:"source"`
|
|
SourceID string `json:"source_id"`
|
|
SplitGroup string `json:"split_group"`
|
|
Tags []string `json:"tags,omitempty"`
|
|
FastPathResolved bool `json:"fast_path_resolved"`
|
|
RouterResidual *bool `json:"router_residual,omitempty"`
|
|
TextHash string `json:"text_hash"`
|
|
Embedding []float32 `json:"embedding"`
|
|
EmbedderID string `json:"embedder_id"`
|
|
FrozenHoldout bool `json:"frozen_holdout"`
|
|
CVFold int `json:"cv_fold"`
|
|
DevPool bool `json:"dev_pool"`
|
|
FamilyID string `json:"family_id"`
|
|
}
|
|
|
|
// ExperimentMeta carries metadata about the experiment run.
|
|
type ExperimentMeta struct {
|
|
EmbedderID string `json:"embedder_id"`
|
|
ModelPath string `json:"model_path"`
|
|
TokenizerPath string `json:"tokenizer_path"`
|
|
Dimension int `json:"dimension"`
|
|
Pooling string `json:"pooling"`
|
|
Normalization string `json:"normalization"`
|
|
InputTemplate string `json:"input_template"`
|
|
TotalExamples int `json:"total_examples"`
|
|
FrozenCount int `json:"frozen_count"`
|
|
DevCount int `json:"dev_count"`
|
|
CVFolds int `json:"cv_folds"`
|
|
FoldComposition map[int]FoldStats `json:"fold_composition"`
|
|
RouteCounts map[string]int `json:"route_counts"`
|
|
SourceCounts map[string]int `json:"source_counts"`
|
|
FastPathCount int `json:"fast_path_count"`
|
|
ResidualCount int `json:"residual_count"`
|
|
HoldoutHash string `json:"holdout_hash"`
|
|
}
|
|
|
|
// FoldStats describes one CV fold.
|
|
type FoldStats struct {
|
|
EvalCount int `json:"eval_count"`
|
|
TrainCount int `json:"train_count"`
|
|
Routes map[string]int `json:"eval_routes"`
|
|
}
|
|
|
|
func main() {
|
|
outPath := flag.String("out", "embeddings.json", "output JSON path")
|
|
folds := flag.Int("folds", 5, "number of CV folds")
|
|
flag.Parse()
|
|
|
|
libPath := os.Getenv("MAVEN_ONNX_LIB")
|
|
if libPath == "" {
|
|
log.Fatal("MAVEN_ONNX_LIB must be set to the libonnxruntime.so path")
|
|
}
|
|
|
|
// Resolve model paths relative to the module root (cwd when running with go run).
|
|
modelPath := "models/embedder/multilingual-e5-small/model_quantized.onnx"
|
|
tokPath := "models/embedder/multilingual-e5-small/tokenizer.json"
|
|
|
|
for _, p := range []string{libPath, modelPath, tokPath} {
|
|
if _, err := os.Stat(p); err != nil {
|
|
log.Fatalf("missing %s: %v", p, err)
|
|
}
|
|
}
|
|
|
|
// Load corpus.
|
|
exs, err := semantic.LoadCorpus()
|
|
if err != nil {
|
|
log.Fatalf("load corpus: %v", err)
|
|
}
|
|
fmt.Fprintf(os.Stderr, "corpus: %d examples\n", len(exs))
|
|
|
|
// Initialize embedder.
|
|
emb, err := router.NewONNXEmbedder(modelPath, tokPath, libPath)
|
|
if err != nil {
|
|
log.Fatalf("init embedder: %v", err)
|
|
}
|
|
defer emb.Close()
|
|
fmt.Fprintf(os.Stderr, "embedder: %s (dim=%d)\n", emb.ID(), emb.Dim())
|
|
|
|
// Compute frozen holdout / dev pool split.
|
|
_, devPool, holdoutHash := semantic.FrozenHoldoutSplit(exs)
|
|
fmt.Fprintf(os.Stderr, "frozen holdout: hash=%s, dev pool: %d examples\n", holdoutHash, len(devPool))
|
|
|
|
// Compute grouped CV folds on dev pool only.
|
|
cvFolds := semantic.GroupedCVFolds(devPool, *folds)
|
|
fmt.Fprintf(os.Stderr, "cv folds: %d\n", len(cvFolds))
|
|
|
|
// Build a lookup: source_id → cv_fold (from dev pool only).
|
|
foldLookup := make(map[string]int)
|
|
for _, f := range cvFolds {
|
|
for _, e := range f.Eval {
|
|
foldLookup[e.SourceID] = f.Fold
|
|
}
|
|
}
|
|
|
|
// Build dev set membership lookup.
|
|
_, devSet, _ := semantic.FrozenHoldoutSplit(exs)
|
|
devIDs := make(map[string]bool)
|
|
for _, e := range devSet {
|
|
devIDs[e.SourceID] = true
|
|
}
|
|
|
|
// Embed all examples.
|
|
ctx := context.Background()
|
|
var cached []CachedRow
|
|
foldComp := make(map[int]*FoldStats)
|
|
for i := 0; i < *folds; i++ {
|
|
foldComp[i] = &FoldStats{Routes: make(map[string]int)}
|
|
}
|
|
routeCounts := make(map[string]int)
|
|
sourceCounts := make(map[string]int)
|
|
fpCount, resCount := 0, 0
|
|
|
|
for i, e := range exs {
|
|
textHash := sha256.Sum256([]byte(e.Text))
|
|
embedding, err := emb.EmbedQuery(ctx, e.Text)
|
|
if err != nil {
|
|
log.Fatalf("embed row %d (%s): %v", i, e.SourceID, err)
|
|
}
|
|
|
|
inDev := devIDs[e.SourceID]
|
|
fold := -1
|
|
if inDev {
|
|
if f, ok := foldLookup[e.SourceID]; ok {
|
|
fold = f
|
|
}
|
|
}
|
|
|
|
isFrozen := !inDev
|
|
|
|
cr := CachedRow{
|
|
Text: e.Text,
|
|
Route: string(e.Route),
|
|
Source: e.Source,
|
|
SourceID: e.SourceID,
|
|
SplitGroup: e.SplitGroup,
|
|
Tags: e.Tags,
|
|
FastPathResolved: e.FastPathResolved,
|
|
RouterResidual: e.RouterResidual,
|
|
TextHash: hex.EncodeToString(textHash[:]),
|
|
Embedding: embedding,
|
|
EmbedderID: emb.ID(),
|
|
FrozenHoldout: isFrozen,
|
|
CVFold: fold,
|
|
DevPool: inDev,
|
|
FamilyID: e.SplitGroup,
|
|
}
|
|
cached = append(cached, cr)
|
|
|
|
routeCounts[cr.Route]++
|
|
sourceCounts[cr.Source]++
|
|
if e.FastPathResolved {
|
|
fpCount++
|
|
} else {
|
|
resCount++
|
|
}
|
|
|
|
if inDev && fold >= 0 {
|
|
foldComp[fold].EvalCount++
|
|
foldComp[fold].Routes[cr.Route]++
|
|
}
|
|
}
|
|
|
|
// Compute train counts per fold.
|
|
for i := 0; i < *folds; i++ {
|
|
foldComp[i].TrainCount = len(devPool) - foldComp[i].EvalCount
|
|
}
|
|
|
|
// Sort route counts for deterministic output.
|
|
sortedRoutes := make([]string, 0, len(routeCounts))
|
|
for r := range routeCounts {
|
|
sortedRoutes = append(sortedRoutes, r)
|
|
}
|
|
sort.Strings(sortedRoutes)
|
|
sortedRouteCounts := make(map[string]int)
|
|
for _, r := range sortedRoutes {
|
|
sortedRouteCounts[r] = routeCounts[r]
|
|
}
|
|
|
|
// Build fold stats with sorted keys.
|
|
finalFoldComp := make(map[int]FoldStats)
|
|
for i := 0; i < *folds; i++ {
|
|
finalFoldComp[i] = *foldComp[i]
|
|
}
|
|
|
|
meta := ExperimentMeta{
|
|
EmbedderID: emb.ID(),
|
|
ModelPath: modelPath,
|
|
TokenizerPath: tokPath,
|
|
Dimension: emb.Dim(),
|
|
Pooling: "mean-pool + L2-normalize",
|
|
Normalization: "L2",
|
|
InputTemplate: "query: <text>",
|
|
TotalExamples: len(exs),
|
|
FrozenCount: len(exs) - len(devPool),
|
|
DevCount: len(devPool),
|
|
CVFolds: *folds,
|
|
FoldComposition: finalFoldComp,
|
|
RouteCounts: sortedRouteCounts,
|
|
SourceCounts: sourceCounts,
|
|
FastPathCount: fpCount,
|
|
ResidualCount: resCount,
|
|
HoldoutHash: holdoutHash,
|
|
}
|
|
|
|
// Output.
|
|
output := map[string]any{
|
|
"meta": meta,
|
|
"examples": cached,
|
|
}
|
|
|
|
data, err := json.MarshalIndent(output, "", " ")
|
|
if err != nil {
|
|
log.Fatalf("marshal: %v", err)
|
|
}
|
|
if err := os.WriteFile(*outPath, data, 0644); err != nil {
|
|
log.Fatalf("write %s: %v", *outPath, err)
|
|
}
|
|
|
|
// Print summary.
|
|
fmt.Fprintf(os.Stderr, "\n=== experiment metadata ===\n")
|
|
fmt.Fprintf(os.Stderr, "embedder: %s\n", meta.EmbedderID)
|
|
fmt.Fprintf(os.Stderr, "dimension: %d\n", meta.Dimension)
|
|
fmt.Fprintf(os.Stderr, "total: %d frozen: %d dev: %d\n", meta.TotalExamples, meta.FrozenCount, meta.DevCount)
|
|
fmt.Fprintf(os.Stderr, "fast-path: %d residual: %d\n", meta.FastPathCount, meta.ResidualCount)
|
|
fmt.Fprintf(os.Stderr, "routes: %s\n", formatMap(sortedRouteCounts))
|
|
fmt.Fprintf(os.Stderr, "fold composition:\n")
|
|
for i := 0; i < *folds; i++ {
|
|
fs := finalFoldComp[i]
|
|
fmt.Fprintf(os.Stderr, " fold %d: eval=%d train=%d routes=%s\n",
|
|
i, fs.EvalCount, fs.TrainCount, formatMap(fs.Routes))
|
|
}
|
|
fmt.Fprintf(os.Stderr, "output: %s (%d bytes)\n", *outPath, len(data))
|
|
}
|
|
|
|
func formatMap(m map[string]int) string {
|
|
var parts []string
|
|
for k, v := range m {
|
|
parts = append(parts, fmt.Sprintf("%s=%d", k, v))
|
|
}
|
|
sort.Strings(parts)
|
|
return "{" + strings.Join(parts, ", ") + "}"
|
|
}
|