302 lines
8.6 KiB
Go
302 lines
8.6 KiB
Go
|
//go:build darwin && arm64
|
||
|
|
|
||
|
|
package ml
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"fmt"
|
||
|
|
"log/slog"
|
||
|
|
"os"
|
||
|
|
"runtime"
|
||
|
|
"sort"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"forge.lthn.ai/core/go-ai/ml"
|
||
|
|
"forge.lthn.ai/core/go/pkg/cli"
|
||
|
|
)
|
||
|
|
|
||
|
|
var benchmarkCmd = &cli.Command{
|
||
|
|
Use: "benchmark",
|
||
|
|
Short: "Compare baseline vs fine-tuned model on ethics probes",
|
||
|
|
Long: `Runs the same prompts through a baseline model and a fine-tuned model,
|
||
|
|
scores both using the heuristic scorer, and outputs a comparison.
|
||
|
|
|
||
|
|
Uses the built-in LEK content probes by default. Optionally takes a
|
||
|
|
custom prompts JSONL file (same format as 'core ml score --input').
|
||
|
|
|
||
|
|
The fine-tuned model can be the same model directory with a LoRA adapter
|
||
|
|
loaded, or a separately merged model.`,
|
||
|
|
RunE: runBenchmark,
|
||
|
|
}
|
||
|
|
|
||
|
|
var (
|
||
|
|
benchmarkBaseline string
|
||
|
|
benchmarkTrained string
|
||
|
|
benchmarkPrompts string
|
||
|
|
benchmarkOutput string
|
||
|
|
benchmarkMaxTokens int
|
||
|
|
benchmarkTemp float64
|
||
|
|
benchmarkMemLimit int
|
||
|
|
)
|
||
|
|
|
||
|
|
func init() {
|
||
|
|
benchmarkCmd.Flags().StringVar(&benchmarkBaseline, "baseline", "", "Path to baseline model directory (required)")
|
||
|
|
benchmarkCmd.Flags().StringVar(&benchmarkTrained, "trained", "", "Path to fine-tuned model directory (required)")
|
||
|
|
benchmarkCmd.Flags().StringVar(&benchmarkPrompts, "prompts", "", "Custom prompts file (JSONL with 'prompt' field, or seeds JSON)")
|
||
|
|
benchmarkCmd.Flags().StringVar(&benchmarkOutput, "output", "benchmark.json", "Output comparison JSON file")
|
||
|
|
benchmarkCmd.Flags().IntVar(&benchmarkMaxTokens, "max-tokens", 1024, "Max tokens per response")
|
||
|
|
benchmarkCmd.Flags().Float64Var(&benchmarkTemp, "temperature", 0.4, "Sampling temperature")
|
||
|
|
benchmarkCmd.Flags().IntVar(&benchmarkMemLimit, "memory-limit", 24, "Metal memory limit in GB")
|
||
|
|
benchmarkCmd.MarkFlagRequired("baseline")
|
||
|
|
benchmarkCmd.MarkFlagRequired("trained")
|
||
|
|
}
|
||
|
|
|
||
|
|
// benchmarkResult holds the comparison for a single prompt.
|
||
|
|
type benchmarkResult struct {
|
||
|
|
ID string `json:"id"`
|
||
|
|
Prompt string `json:"prompt"`
|
||
|
|
BaselineResponse string `json:"baseline_response"`
|
||
|
|
TrainedResponse string `json:"trained_response"`
|
||
|
|
BaselineLEK float64 `json:"baseline_lek_score"`
|
||
|
|
TrainedLEK float64 `json:"trained_lek_score"`
|
||
|
|
Delta float64 `json:"delta"`
|
||
|
|
|
||
|
|
BaselineHeuristic *ml.HeuristicScores `json:"baseline_heuristic"`
|
||
|
|
TrainedHeuristic *ml.HeuristicScores `json:"trained_heuristic"`
|
||
|
|
}
|
||
|
|
|
||
|
|
// benchmarkSummary holds aggregate comparison metrics.
|
||
|
|
type benchmarkSummary struct {
|
||
|
|
BaselineModel string `json:"baseline_model"`
|
||
|
|
TrainedModel string `json:"trained_model"`
|
||
|
|
TotalPrompts int `json:"total_prompts"`
|
||
|
|
AvgBaselineLEK float64 `json:"avg_baseline_lek"`
|
||
|
|
AvgTrainedLEK float64 `json:"avg_trained_lek"`
|
||
|
|
AvgDelta float64 `json:"avg_delta"`
|
||
|
|
Improved int `json:"improved"`
|
||
|
|
Regressed int `json:"regressed"`
|
||
|
|
Unchanged int `json:"unchanged"`
|
||
|
|
Duration string `json:"duration"`
|
||
|
|
Results []benchmarkResult `json:"results"`
|
||
|
|
}
|
||
|
|
|
||
|
|
func runBenchmark(cmd *cli.Command, args []string) error {
|
||
|
|
start := time.Now()
|
||
|
|
|
||
|
|
// Load prompts — either custom file or built-in probes
|
||
|
|
prompts, err := loadBenchmarkPrompts()
|
||
|
|
if err != nil {
|
||
|
|
return err
|
||
|
|
}
|
||
|
|
|
||
|
|
slog.Info("benchmark: loaded prompts", "count", len(prompts))
|
||
|
|
|
||
|
|
opts := ml.GenOpts{
|
||
|
|
Temperature: benchmarkTemp,
|
||
|
|
MaxTokens: benchmarkMaxTokens,
|
||
|
|
}
|
||
|
|
|
||
|
|
// Generate baseline responses
|
||
|
|
slog.Info("benchmark: loading baseline model", "path", benchmarkBaseline)
|
||
|
|
baselineBackend, err := ml.NewMLXBackend(benchmarkBaseline)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("load baseline: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
baselineResponses := make(map[string]string)
|
||
|
|
for i, p := range prompts {
|
||
|
|
slog.Info("benchmark: baseline",
|
||
|
|
"prompt", fmt.Sprintf("%d/%d", i+1, len(prompts)),
|
||
|
|
"id", p.id,
|
||
|
|
)
|
||
|
|
resp, err := baselineBackend.Generate(context.Background(), p.prompt, opts)
|
||
|
|
if err != nil {
|
||
|
|
slog.Error("benchmark: baseline failed", "id", p.id, "error", err)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
baselineResponses[p.id] = resp
|
||
|
|
|
||
|
|
if (i+1)%4 == 0 {
|
||
|
|
runtime.GC()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Force cleanup before loading second model
|
||
|
|
baselineBackend = nil
|
||
|
|
runtime.GC()
|
||
|
|
runtime.GC()
|
||
|
|
|
||
|
|
// Generate trained responses
|
||
|
|
slog.Info("benchmark: loading trained model", "path", benchmarkTrained)
|
||
|
|
trainedBackend, err := ml.NewMLXBackend(benchmarkTrained)
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("load trained: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
trainedResponses := make(map[string]string)
|
||
|
|
for i, p := range prompts {
|
||
|
|
slog.Info("benchmark: trained",
|
||
|
|
"prompt", fmt.Sprintf("%d/%d", i+1, len(prompts)),
|
||
|
|
"id", p.id,
|
||
|
|
)
|
||
|
|
resp, err := trainedBackend.Generate(context.Background(), p.prompt, opts)
|
||
|
|
if err != nil {
|
||
|
|
slog.Error("benchmark: trained failed", "id", p.id, "error", err)
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
trainedResponses[p.id] = resp
|
||
|
|
|
||
|
|
if (i+1)%4 == 0 {
|
||
|
|
runtime.GC()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
trainedBackend = nil
|
||
|
|
runtime.GC()
|
||
|
|
|
||
|
|
// Score both sets
|
||
|
|
var results []benchmarkResult
|
||
|
|
var totalBaseline, totalTrained float64
|
||
|
|
improved, regressed, unchanged := 0, 0, 0
|
||
|
|
|
||
|
|
for _, p := range prompts {
|
||
|
|
baseResp := baselineResponses[p.id]
|
||
|
|
trainResp := trainedResponses[p.id]
|
||
|
|
|
||
|
|
if baseResp == "" || trainResp == "" {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
|
||
|
|
baseH := ml.ScoreHeuristic(baseResp)
|
||
|
|
trainH := ml.ScoreHeuristic(trainResp)
|
||
|
|
delta := trainH.LEKScore - baseH.LEKScore
|
||
|
|
|
||
|
|
totalBaseline += baseH.LEKScore
|
||
|
|
totalTrained += trainH.LEKScore
|
||
|
|
|
||
|
|
if delta > 0.5 {
|
||
|
|
improved++
|
||
|
|
} else if delta < -0.5 {
|
||
|
|
regressed++
|
||
|
|
} else {
|
||
|
|
unchanged++
|
||
|
|
}
|
||
|
|
|
||
|
|
results = append(results, benchmarkResult{
|
||
|
|
ID: p.id,
|
||
|
|
Prompt: p.prompt,
|
||
|
|
BaselineResponse: baseResp,
|
||
|
|
TrainedResponse: trainResp,
|
||
|
|
BaselineLEK: baseH.LEKScore,
|
||
|
|
TrainedLEK: trainH.LEKScore,
|
||
|
|
Delta: delta,
|
||
|
|
BaselineHeuristic: baseH,
|
||
|
|
TrainedHeuristic: trainH,
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
n := float64(len(results))
|
||
|
|
if n == 0 {
|
||
|
|
return fmt.Errorf("no results to compare")
|
||
|
|
}
|
||
|
|
|
||
|
|
summary := benchmarkSummary{
|
||
|
|
BaselineModel: benchmarkBaseline,
|
||
|
|
TrainedModel: benchmarkTrained,
|
||
|
|
TotalPrompts: len(results),
|
||
|
|
AvgBaselineLEK: totalBaseline / n,
|
||
|
|
AvgTrainedLEK: totalTrained / n,
|
||
|
|
AvgDelta: (totalTrained - totalBaseline) / n,
|
||
|
|
Improved: improved,
|
||
|
|
Regressed: regressed,
|
||
|
|
Unchanged: unchanged,
|
||
|
|
Duration: time.Since(start).Round(time.Second).String(),
|
||
|
|
Results: results,
|
||
|
|
}
|
||
|
|
|
||
|
|
// Write output
|
||
|
|
data, err := json.MarshalIndent(summary, "", " ")
|
||
|
|
if err != nil {
|
||
|
|
return fmt.Errorf("marshal output: %w", err)
|
||
|
|
}
|
||
|
|
if err := os.WriteFile(benchmarkOutput, data, 0644); err != nil {
|
||
|
|
return fmt.Errorf("write output: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Print summary
|
||
|
|
fmt.Println()
|
||
|
|
fmt.Println("=== Benchmark Results ===")
|
||
|
|
fmt.Printf("Baseline: %s\n", benchmarkBaseline)
|
||
|
|
fmt.Printf("Trained: %s\n", benchmarkTrained)
|
||
|
|
fmt.Printf("Prompts: %d\n", len(results))
|
||
|
|
fmt.Println()
|
||
|
|
fmt.Printf("Avg LEK (baseline): %+.2f\n", summary.AvgBaselineLEK)
|
||
|
|
fmt.Printf("Avg LEK (trained): %+.2f\n", summary.AvgTrainedLEK)
|
||
|
|
fmt.Printf("Avg Delta: %+.2f\n", summary.AvgDelta)
|
||
|
|
fmt.Println()
|
||
|
|
fmt.Printf("Improved: %d (%.0f%%)\n", improved, float64(improved)/n*100)
|
||
|
|
fmt.Printf("Regressed: %d (%.0f%%)\n", regressed, float64(regressed)/n*100)
|
||
|
|
fmt.Printf("Unchanged: %d (%.0f%%)\n", unchanged, float64(unchanged)/n*100)
|
||
|
|
fmt.Printf("Duration: %s\n", summary.Duration)
|
||
|
|
fmt.Printf("Output: %s\n", benchmarkOutput)
|
||
|
|
|
||
|
|
return nil
|
||
|
|
}
|
||
|
|
|
||
|
|
type benchPrompt struct {
|
||
|
|
id string
|
||
|
|
prompt string
|
||
|
|
}
|
||
|
|
|
||
|
|
func loadBenchmarkPrompts() ([]benchPrompt, error) {
|
||
|
|
if benchmarkPrompts == "" {
|
||
|
|
// Use built-in content probes
|
||
|
|
probes := ml.ContentProbes
|
||
|
|
prompts := make([]benchPrompt, len(probes))
|
||
|
|
for i, p := range probes {
|
||
|
|
prompts[i] = benchPrompt{id: p.ID, prompt: p.Prompt}
|
||
|
|
}
|
||
|
|
return prompts, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Try seeds JSON format first (array of {id, prompt, ...})
|
||
|
|
data, err := os.ReadFile(benchmarkPrompts)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("read prompts: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
var seeds []seedPrompt
|
||
|
|
if json.Unmarshal(data, &seeds) == nil && len(seeds) > 0 {
|
||
|
|
prompts := make([]benchPrompt, len(seeds))
|
||
|
|
for i, s := range seeds {
|
||
|
|
prompts[i] = benchPrompt{id: s.ID, prompt: s.Prompt}
|
||
|
|
}
|
||
|
|
return prompts, nil
|
||
|
|
}
|
||
|
|
|
||
|
|
// Try JSONL responses format
|
||
|
|
responses, err := ml.ReadResponses(benchmarkPrompts)
|
||
|
|
if err != nil {
|
||
|
|
return nil, fmt.Errorf("parse prompts: %w", err)
|
||
|
|
}
|
||
|
|
|
||
|
|
// Deduplicate by prompt
|
||
|
|
seen := make(map[string]bool)
|
||
|
|
var prompts []benchPrompt
|
||
|
|
for _, r := range responses {
|
||
|
|
if seen[r.Prompt] {
|
||
|
|
continue
|
||
|
|
}
|
||
|
|
seen[r.Prompt] = true
|
||
|
|
id := r.ID
|
||
|
|
if id == "" {
|
||
|
|
id = fmt.Sprintf("P%03d", len(prompts)+1)
|
||
|
|
}
|
||
|
|
prompts = append(prompts, benchPrompt{id: id, prompt: r.Prompt})
|
||
|
|
}
|
||
|
|
|
||
|
|
sort.Slice(prompts, func(i, j int) bool { return prompts[i].id < prompts[j].id })
|
||
|
|
return prompts, nil
|
||
|
|
}
|