diff --git a/cmd/ml/cmd_sandwich.go b/cmd/ml/cmd_sandwich.go new file mode 100644 index 00000000..9fcad6d9 --- /dev/null +++ b/cmd/ml/cmd_sandwich.go @@ -0,0 +1,238 @@ +//go:build darwin && arm64 + +package ml + +import ( + "context" + "encoding/json" + "fmt" + "log/slog" + "os" + "runtime" + "time" + + "forge.lthn.ai/core/go-ai/ml" + "forge.lthn.ai/core/go/pkg/cli" +) + +var sandwichCmd = &cli.Command{ + Use: "sandwich", + Short: "Generate LEK training data using sandwich signing", + Long: `Generates training data by wrapping seed prompts in a "sandwich" format: + + KB preamble (axioms framework) → seed prompt → LEK-1 kernel postfix + +Each seed prompt is sent to the local MLX model for inference, and the +signed prompt + response pair is written as chat JSONL for 'core ml train'. + +The "sandwich" format embeds the ethical framework context around each +prompt, teaching the model to reason from LEK principles naturally. + +Seed file format (JSON array): + [{"id": "P01", "category": "sovereignty", "prompt": "...", "signal": "..."}]`, + RunE: runSandwich, +} + +var ( + sandwichModelPath string + sandwichKB string + sandwichKernel string + sandwichSeeds string + sandwichOutput string + sandwichMaxTokens int + sandwichTemp float64 + sandwichMemLimit int + sandwichDryRun bool +) + +func init() { + sandwichCmd.Flags().StringVar(&sandwichModelPath, "model-path", "", "Path to model directory (required)") + sandwichCmd.Flags().StringVar(&sandwichKB, "kb", "", "Knowledge base document (axioms markdown, required)") + sandwichCmd.Flags().StringVar(&sandwichKernel, "kernel", "", "LEK-1 kernel file (required)") + sandwichCmd.Flags().StringVar(&sandwichSeeds, "seeds", "", "Seed prompts JSON file (required)") + sandwichCmd.Flags().StringVar(&sandwichOutput, "output", "sandwich.jsonl", "Output JSONL file") + sandwichCmd.Flags().IntVar(&sandwichMaxTokens, "max-tokens", 1024, "Max tokens per response") + sandwichCmd.Flags().Float64Var(&sandwichTemp, "temperature", 0.4, "Sampling temperature") + sandwichCmd.Flags().IntVar(&sandwichMemLimit, "memory-limit", 24, "Metal memory limit in GB") + sandwichCmd.Flags().BoolVar(&sandwichDryRun, "dry-run", false, "Output prompts only (no inference)") + sandwichCmd.MarkFlagRequired("model-path") + sandwichCmd.MarkFlagRequired("kernel") + sandwichCmd.MarkFlagRequired("seeds") + sandwichCmd.MarkFlagRequired("kb") +} + +// seedPrompt is a single prompt from the seeds JSON file. +type seedPrompt struct { + ID string `json:"id"` + Category string `json:"category"` + Prompt string `json:"prompt"` + Signal string `json:"signal"` +} + +// sandwichOutput holds a single training example in messages format. +type sandwichRecord struct { + Messages []ml.Message `json:"messages"` +} + +func runSandwich(cmd *cli.Command, args []string) error { + start := time.Now() + + // Load KB document + kbBytes, err := os.ReadFile(sandwichKB) + if err != nil { + return fmt.Errorf("read KB: %w", err) + } + kbText := string(kbBytes) + + // Load LEK-1 kernel + kernelBytes, err := os.ReadFile(sandwichKernel) + if err != nil { + return fmt.Errorf("read kernel: %w", err) + } + kernelText := string(kernelBytes) + + // Load seed prompts + seedBytes, err := os.ReadFile(sandwichSeeds) + if err != nil { + return fmt.Errorf("read seeds: %w", err) + } + var seeds []seedPrompt + if err := json.Unmarshal(seedBytes, &seeds); err != nil { + return fmt.Errorf("parse seeds: %w", err) + } + + slog.Info("sandwich: loaded inputs", + "kb_chars", len(kbText), + "kernel_chars", len(kernelText), + "seeds", len(seeds), + ) + + if len(seeds) == 0 { + return fmt.Errorf("no seed prompts found") + } + + // Open output file + outFile, err := os.Create(sandwichOutput) + if err != nil { + return fmt.Errorf("create output: %w", err) + } + defer outFile.Close() + encoder := json.NewEncoder(outFile) + + // Dry-run mode: output prompts without inference + if sandwichDryRun { + for _, seed := range seeds { + signedPrompt := buildSandwich(kbText, seed.Prompt, kernelText) + record := sandwichRecord{ + Messages: []ml.Message{ + {Role: "user", Content: signedPrompt}, + }, + } + if err := encoder.Encode(record); err != nil { + return fmt.Errorf("write record: %w", err) + } + } + slog.Info("sandwich: dry-run complete", + "output", sandwichOutput, + "prompts", len(seeds), + ) + return nil + } + + // Load MLX model + slog.Info("sandwich: loading model", "path", sandwichModelPath) + backend, err := ml.NewMLXBackend(sandwichModelPath) + if err != nil { + return fmt.Errorf("load model: %w", err) + } + + opts := ml.GenOpts{ + Temperature: sandwichTemp, + MaxTokens: sandwichMaxTokens, + } + + var totalTokenTime time.Duration + generated := 0 + + for i, seed := range seeds { + seedStart := time.Now() + + // Build the sandwich: KB + prompt + kernel + signedPrompt := buildSandwich(kbText, seed.Prompt, kernelText) + + // Send as a user message for chat-style generation + messages := []ml.Message{ + {Role: "user", Content: signedPrompt}, + } + + slog.Info("sandwich: generating", + "seed", fmt.Sprintf("%d/%d", i+1, len(seeds)), + "id", seed.ID, + "category", seed.Category, + ) + + // Generate response + response, err := backend.Chat(context.Background(), messages, opts) + if err != nil { + slog.Error("sandwich: generation failed", + "id", seed.ID, + "error", err, + ) + continue + } + + elapsed := time.Since(seedStart) + totalTokenTime += elapsed + + // Write training record + record := sandwichRecord{ + Messages: []ml.Message{ + {Role: "user", Content: signedPrompt}, + {Role: "assistant", Content: response}, + }, + } + if err := encoder.Encode(record); err != nil { + return fmt.Errorf("write record: %w", err) + } + + generated++ + slog.Info("sandwich: generated", + "id", seed.ID, + "category", seed.Category, + "response_chars", len(response), + "duration", elapsed.Round(time.Second), + ) + + // Periodic cleanup + if (i+1)%4 == 0 { + runtime.GC() + } + } + + slog.Info("sandwich: complete", + "output", sandwichOutput, + "generated", generated, + "total", len(seeds), + "duration", time.Since(start).Round(time.Second), + "avg_per_seed", (totalTokenTime / time.Duration(max(generated, 1))).Round(time.Second), + ) + + return nil +} + +// buildSandwich constructs the signed prompt: KB preamble + seed prompt + LEK-1 kernel. +func buildSandwich(kb, prompt, kernel string) string { + return fmt.Sprintf(`Name: Ethics Experiment +KB: +%s + +--- + +%s + +--- + +%s + +Remember: respond using the ethical framework above. Do not reference the framework directly — reason from its principles naturally.`, kb, prompt, kernel) +} diff --git a/cmd/ml/cmd_sandwich_init.go b/cmd/ml/cmd_sandwich_init.go new file mode 100644 index 00000000..d23b8856 --- /dev/null +++ b/cmd/ml/cmd_sandwich_init.go @@ -0,0 +1,7 @@ +//go:build darwin && arm64 + +package ml + +func init() { + mlCmd.AddCommand(sandwichCmd) +}