diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,26 @@ ## Unreleased +## 0.2.1.0 - 2026-07-05 + +### Added + +- Shared optimizer budget metering helpers, effective-instruction inspection, and + exact LLM-call counting for opaque optimizer runs. +- Budgeted ensemble search via `ensembleSearchWith` and exported + `bootstrapKeptDemos` for metered bootstrap demo recovery. + +### Changed + +- Optimizers now reserve predicted LM-call cost before scoring/proposing and + return the best candidate found so far when the next spend would exceed budget. +- Instruction search and seeding preserve effective instructions without + serializing redundant overrides. +- Ensemble and KNN documentation now calls out structure-changing artifacts and + how to load saved state against matching templates. +- Refreshed internal `shikumi`, `shikumi-compile`, `shikumi-eval`, and + `shikumi-trace` bounds for the current package set. + ## 0.2.0.0 - 2026-06-28 ### Added diff --git a/shikumi-optimize.cabal b/shikumi-optimize.cabal --- a/shikumi-optimize.cabal +++ b/shikumi-optimize.cabal @@ -1,6 +1,6 @@ cabal-version: 3.4 name: shikumi-optimize -version: 0.2.0.0 +version: 0.2.1.0 synopsis: The optimizer framework for shikumi LM programs (EP-10) category: AI description: @@ -66,10 +66,10 @@ , effectful , generic-lens ^>=2.2 , lens ^>=5.3 - , shikumi ^>=0.2.0.0 - , shikumi-compile ^>=0.1.1.0 - , shikumi-eval ^>=0.1.1.0 - , shikumi-trace ^>=0.1.1.0 + , shikumi ^>=0.3.0.0 + , shikumi-compile ^>=0.2.0.0 + , shikumi-eval ^>=0.2.0.0 + , shikumi-trace ^>=0.2.0.0 , text ^>=2.1 , vector @@ -92,21 +92,23 @@ OptimizeSpec ProposeSpec RandomSearchSpec + SearchSpec + SeedingSpec StubLM build-depends: , aeson - , baikai >=0.2 && <0.3 + , baikai >=0.3 && <0.4 , base , containers , effectful , generic-lens , lens - , shikumi ^>=0.2.0.0 - , shikumi-compile ^>=0.1.1.0 - , shikumi-eval ^>=0.1.1.0 - , shikumi-optimize ^>=0.2.0.0 - , shikumi-trace ^>=0.1.1.0 + , shikumi ^>=0.3.0.0 + , shikumi-compile ^>=0.2.0.0 + , shikumi-eval ^>=0.2.0.0 + , shikumi-optimize ^>=0.2.1.0 + , shikumi-trace ^>=0.2.0.0 , tasty , tasty-hunit , text diff --git a/src/Shikumi/Optimize/Bootstrap.hs b/src/Shikumi/Optimize/Bootstrap.hs --- a/src/Shikumi/Optimize/Bootstrap.hs +++ b/src/Shikumi/Optimize/Bootstrap.hs @@ -21,6 +21,7 @@ module Shikumi.Optimize.Bootstrap ( bootstrapFewShot, bootstrapFewShotWith, + bootstrapKeptDemos, BootstrapConfig (..), defaultBootstrapConfig, recoverDemo, @@ -28,14 +29,17 @@ where import Data.Aeson (ToJSON, toJSON) -import Effectful.Error.Static (catchError) +import Effectful (Eff, (:>)) +import Effectful.Error.Static (Error, catchError) +import Effectful.Prim (Prim) import GHC.Generics (Generic) import Shikumi.Error (ShikumiError) -import Shikumi.Eval (Example (..), datasetExamples, prediction, unScore) +import Shikumi.Eval (Dataset, Example (..), Metric, datasetExamples, prediction, unScore) +import Shikumi.LLM (LLM) import Shikumi.Optimize.LabeledFewShot (withDemos) -import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Search (BudgetMeter, freezeProgram, newBudgetMeter, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) -import Shikumi.Program (Demo (..), Program, runProgram) +import Shikumi.Program (Demo (..), Program, foldParams, runProgram) -- | Tunables for a bootstrap search. data BootstrapConfig = BootstrapConfig @@ -64,7 +68,10 @@ -- | Bootstrap few-shot with an explicit configuration. The @teacher@ may be a -- stronger or chain-of-thought variant of the student, or the student itself; it --- must share the student's input/output types. +-- must share the student's input/output types. Each teacher run reserves one +-- predicted LM completion per teacher predict node before it runs; when the next +-- teacher run does not fit the 'Budget', demo recovery stops and the demos found so +-- far are attached. bootstrapFewShotWith :: (ToJSON i, ToJSON o) => BootstrapConfig -> @@ -73,11 +80,22 @@ Budget -> Optimizer i o bootstrapFewShotWith cfg teacher budget = Optimizer $ \train metric student -> do - -- Bound the number of teacher runs by the budget (each run is >= 1 LM call). - let exs = take (max 0 (maxLmCalls budget)) (datasetExamples train) - -- Run the teacher on one example; if it succeeds and the metric passes, - -- emit its recovered demo, else emit nothing. Recovery is total: a teacher - -- error yields no demo rather than aborting the search. + meter <- newBudgetMeter budget + kept <- bootstrapKeptDemos cfg meter teacher train metric + pure (freezeProgram (withDemos kept student)) + +-- | Recover metric-passing demos from teacher runs under a shared budget meter. +bootstrapKeptDemos :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => + BootstrapConfig -> + BudgetMeter -> + Program i o -> + Dataset i o -> + Metric o -> + Eff es [Demo] +bootstrapKeptDemos cfg meter teacher train metric = do + let teacherCost = max 1 (length (foldParams teacher)) + cap = max 0 (maxBootstrappedDemos cfg) keepIfPassing (Example inp expd) = ( do out <- runProgram teacher inp @@ -85,5 +103,16 @@ pure [recoverDemo inp out | s >= passThreshold cfg] ) `catchError` \_ (_ :: ShikumiError) -> pure [] - kept <- concat <$> mapM keepIfPassing exs - pure (freezeProgram (withDemos (take (maxBootstrappedDemos cfg) kept) student)) + collect kept [] + | length kept >= cap = pure (take cap kept) + | otherwise = pure kept + collect kept (ex : rest) + | length kept >= cap = pure (take cap kept) + | otherwise = do + fits <- tryCharge meter teacherCost + if not fits + then pure kept + else do + newKept <- keepIfPassing ex + collect (kept ++ newKept) rest + collect [] (datasetExamples train) diff --git a/src/Shikumi/Optimize/COPRO.hs b/src/Shikumi/Optimize/COPRO.hs --- a/src/Shikumi/Optimize/COPRO.hs +++ b/src/Shikumi/Optimize/COPRO.hs @@ -9,7 +9,8 @@ -- COPRO consumes EP-19's grounded proposer ('Shikumi.Optimize.Propose.proposeInstructions') -- directly: each round's call passes the node's current instruction and its scored -- 'PastInstruction' history, and the proposer returns ranked candidates with the --- current instruction always retained (the safety property — a node never degrades). +-- current effective instruction always retained. Keeping that candidate writes no +-- redundant override, preserving the safety property that a node never degrades. -- -- Output is V1's 'Shikumi.Compile.Types.CompiledProgram' via 'freezeProgram', invoked -- through 'Shikumi.Optimize.optimize' and serialized unchanged (integration point #4). @@ -22,7 +23,6 @@ import Control.Monad (foldM) import Data.Aeson (ToJSON) -import Data.Maybe (fromMaybe) import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent) import Effectful.Error.Static (Error) @@ -30,7 +30,7 @@ import GHC.Generics (Generic) import Shikumi.Effect.Time (Time) import Shikumi.Error (ShikumiError) -import Shikumi.Eval (Dataset, Metric, datasetSize) +import Shikumi.Eval (Dataset, Metric) import Shikumi.LLM (LLM) import Shikumi.Optimize.Propose ( PastInstruction (..), @@ -38,7 +38,7 @@ ProposeResult (..), proposeInstructions, ) -import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr) +import Shikumi.Optimize.Search (BudgetMeter, effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) import Shikumi.Program (Program, foldParams) @@ -59,19 +59,22 @@ -- | Coordinate-ascent instruction optimization. Visits each node in @foldParams@ -- order, optimizing it over @depth@ rounds against the already-improved earlier --- nodes, threading one running LM-call count so the 'Budget' bounds the whole search. +-- nodes. Proposer calls and candidate scoring reserve their predicted cost through +-- one shared 'Budget', so the search returns the best-so-far before the next spend +-- would exceed either ceiling. copro :: (ToJSON i, ToJSON o) => CoproConfig -> Optimizer i o copro cfg = Optimizer $ \train metric student -> do + meter <- newBudgetMeter (budget cfg) let nNodes = length (foldParams student) - (final, _calls) <- + final <- foldM - (\acc idx -> optimizeNode cfg train metric idx acc) - (student, 0) + (\acc idx -> optimizeNode cfg meter train metric idx acc) + student [0 .. nNodes - 1] pure (freezeProgram final) -- | Optimize node @idx@ over @depth@ rounds. Returns the program with node @idx@ set --- to its best-found instruction, plus the updated LM-call count. Each round proposes +-- to its best-found instruction. Each round proposes -- @breadth - 1@ fresh candidates (plus the retained current instruction) via the -- grounded proposer fed the scored attempt history, scores the not-yet-seen ones on -- the whole training set, records @(instruction, best-score)@, and sets the node to @@ -79,27 +82,26 @@ optimizeNode :: (ToJSON i, ToJSON o, LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => CoproConfig -> + BudgetMeter -> Dataset i o -> Metric o -> Int -> - (Program i o, Int) -> - Eff es (Program i o, Int) -optimizeNode cfg train metric idx (prog0, calls0) = goRound 1 prog0 calls0 [] + Program i o -> + Eff es (Program i o) +optimizeNode cfg meter train metric idx prog0 = goRound 1 prog0 [] where - dsSize = datasetSize train bdth = max 2 (breadth cfg) dpth = max 1 (depth cfg) proposerCost = 4 + (bdth - 1) - maxCalls = maxLmCalls (budget cfg) - maxCands = maxCandidates (budget cfg) - goRound r prog calls evald - | r > dpth = pure (setBest prog evald, calls) + goRound r prog evald + | r > dpth = pure (setBest prog evald) | otherwise = do - let cur = fromMaybe "" (instructionAt idx prog) + let cur = effectiveInstructionAt idx prog hist = [PastInstruction i s | (i, s) <- evald] - (cands, calls1) <- - if calls + proposerCost <= maxCalls + fitsProposer <- tryCharge meter proposerCost + cands <- + if fitsProposer then do ProposeResult cs <- proposeInstructions @@ -114,22 +116,21 @@ tipIndex = 1, viewBatch = 2 } - pure (cs, calls + proposerCost) - else pure ([cur], calls) - (evald', calls2) <- scoreNew prog calls1 evald (dedupNew evald cands) + pure cs + else pure [cur] + evald' <- scoreNew prog evald (dedupNew evald cands) -- set the node to the best instruction found so far, then continue let prog' = setBest prog evald' - goRound (r + 1) prog' calls2 evald' + goRound (r + 1) prog' evald' - -- Score the not-yet-evaluated candidates, threading the call count and stopping + -- Score the not-yet-evaluated candidates, using the shared meter and stopping -- before either Budget ceiling would be exceeded. - scoreNew _ calls evald [] = pure (evald, calls) - scoreNew prog calls evald (c : cs) - | calls + dsSize > maxCalls = pure (evald, calls) - | length evald >= maxCands = pure (evald, calls) - | otherwise = do - s <- scoreOn train metric (setNodeInstr idx c prog) - scoreNew prog (calls + dsSize) (evald ++ [(c, s)]) cs + scoreNew _ evald [] = pure evald + scoreNew prog evald (c : cs) = do + ms <- meteredScore meter train metric (setNodeInstrIfNew idx c prog) + case ms of + Nothing -> pure evald + Just s -> scoreNew prog (evald ++ [(c, s)]) cs -- Candidates not already scored, de-duplicated against each other (order-preserving). dedupNew evald = go (map fst evald) @@ -142,6 +143,6 @@ -- Set node idx to the highest-scoring instruction recorded (earliest on ties). setBest prog evald = case evald of [] -> prog - (e : es) -> setNodeInstr idx (fst (foldl' pick e es)) prog + (e : es) -> setNodeInstrIfNew idx (fst (foldl' pick e es)) prog where pick best c = if snd c > snd best then c else best diff --git a/src/Shikumi/Optimize/Ensemble.hs b/src/Shikumi/Optimize/Ensemble.hs --- a/src/Shikumi/Optimize/Ensemble.hs +++ b/src/Shikumi/Optimize/Ensemble.hs @@ -9,8 +9,16 @@ -- -- Resampling is deterministic (a fixed linear-congruential stream seeded by the -- member index), so the search is reproducible run to run. +-- +-- The returned artifact is structure-changing: it is an @Ensemble@ combinator over +-- the optimized members, not just a parameter rewrite of the student. The member +-- parameters are persisted by 'Shikumi.Compile.Serialize.encodeCompiled', but the +-- reducer closure and exact member structure live in the program template held in +-- code. Load saved state onto the matching ensemble template, not onto the plain +-- student. module Shikumi.Optimize.Ensemble ( ensembleSearch, + ensembleSearchWith, majorityReducer, ) where @@ -19,18 +27,29 @@ import Shikumi.Combinator (ensemble) import Shikumi.Compile.Types (compiledProgram) import Shikumi.Eval (Dataset, dataset, datasetExamples) -import Shikumi.Optimize.Search (freezeProgram) -import Shikumi.Optimize.Types (Optimizer (..)) +import Shikumi.Optimize.Search (freezeProgram, withLmCallCount) +import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) -- | Build an @size@-member ensemble: run @inner@ on @size@ bootstrap resamples of --- the training set and combine the resulting programs by majority vote. +-- the training set and combine the resulting programs by majority vote. This +-- changes structure by returning an @Ensemble@; saved state must be decoded onto a +-- matching ensemble template. ensembleSearch :: (Eq o) => Int -> Optimizer i o -> Optimizer i o -ensembleSearch size inner = Optimizer $ \train metric student -> do +ensembleSearch = ensembleSearchWith defaultBudget + +-- | Build a budgeted ensemble. The budget is enforced between members using exact +-- LLM-call counting; at least one member always runs, so the final spend may exceed +-- the bound by one member's cost. +ensembleSearchWith :: (Eq o) => Budget -> Int -> Optimizer i o -> Optimizer i o +ensembleSearchWith budget size inner = Optimizer $ \train metric student -> do let seeds = [1 .. max 1 size] - members <- - mapM - (\seed -> compiledProgram <$> runOptimizer inner (resample seed train) metric student) - seeds + go _total acc [] = pure (reverse acc) + go total acc (seed : rest) + | total >= maxLmCalls budget && not (null acc) = pure (reverse acc) + | otherwise = do + (cp, calls) <- withLmCallCount (runOptimizer inner (resample seed train) metric student) + go (total + calls) (compiledProgram cp : acc) rest + members <- go 0 [] seeds pure (freezeProgram (ensemble members majorityReducer)) -- | Sample a dataset with replacement, deterministically, seeded by @seed@. An diff --git a/src/Shikumi/Optimize/GEPA.hs b/src/Shikumi/Optimize/GEPA.hs --- a/src/Shikumi/Optimize/GEPA.hs +++ b/src/Shikumi/Optimize/GEPA.hs @@ -33,10 +33,7 @@ ) where -import Control.Lens ((&), (?~)) import Control.Monad (forM, forM_, when) -import Data.Generics.Labels () -import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text qualified as T import Effectful (Eff, (:>)) @@ -61,14 +58,12 @@ import Shikumi.LLM (LLM) import Shikumi.Module (predict) import Shikumi.Optimize.Pareto (Candidate (..), paretoFrontier, sampleParent) -import Shikumi.Optimize.Search (freezeProgram) +import Shikumi.Optimize.Search (effectiveInstructionAt, freezeProgram, newBudgetMeter, scoringCost, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) import Shikumi.Program ( NodeFields (..), - Params (..), Program, foldParams, - mapParamsAt, nodeFieldsIndexed, runProgram, setProgramParams, @@ -176,18 +171,12 @@ in if null crits then pure prog else do - let cur = fromMaybe "" (instructionOverride (paramsAt idx prog)) + let cur = effectiveInstructionAt idx prog fldSummary = renderFields (drop idx fields) fb = T.intercalate "\n" crits ReflectOut newInstr <- runProgram proposer (ReflectIn cur fb progSummary dataSummary fldSummary) - pure (mapParamsAt idx (\ps -> ps & #instructionOverride ?~ newInstr) prog) - --- | The 'Params' at a node index (empty if out of range). -paramsAt :: Int -> Program i o -> Params -paramsAt idx prog = case drop idx (foldParams prog) of - (ps : _) -> ps - [] -> Params Nothing [] + pure (setNodeInstrIfNew idx newInstr prog) -- | Render a node's field names for the proposer prompt. renderFields :: [NodeFields] -> Text @@ -203,54 +192,66 @@ -- | The reflective evolutionary optimizer. Takes its reflective proposer and feedback -- metric explicitly (so it is testable under a stub LM) and returns V1's --- 'Optimizer'. +-- 'Optimizer'. GEPA gates its seed evaluation before any LM call; if the budget is +-- too small to score the student once, it returns the student unscored. Each +-- evolution step reserves a conservative full-step cost before capture, reflection, +-- and child scoring. gepa :: Program ReflectIn ReflectOut -> FeedbackMetric o -> Budget -> Optimizer i o gepa proposer fbMetric budget = Optimizer $ \train metric student -> do + meter <- newBudgetMeter budget let paths = programNodePaths student fields = nodeFieldsIndexed student nNodes = max 1 (length paths) - n = max 1 (datasetSize train) progSummary = fallbackProgramSummary (length paths) dataSummary = fallbackDatasetSummary (datasetSize train) - maxCalls = maxLmCalls budget maxCands = maxCandidates budget rebuild cand = either (const student) id (setProgramParams (params cand) student) + seedCost = scoringCost train student - seedRpt <- evaluatePure train metric student - let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt) + seedFits <- tryCharge meter seedCost + if not seedFits + then pure (freezeProgram student) + else do + seedRpt <- evaluatePure train metric student + let seedCand = Candidate (foldParams student) (perEx seedRpt) (aggregateScore seedRpt) - -- A full step costs: capture (n) + mutate (1) + child eval (n) = 2n + 1. - stepCost = 2 * n + 1 - stepCap = maxCands + 4 + -- A full step costs: capture + child evaluation over the whole dataset, + -- plus one reflective proposer call. + stepCost = 2 * seedCost + 1 + stepCap = maxCands + 4 - loop step calls cands seed frontier - | step >= stepCap = pure (bestOf seedCand frontier) - | calls + stepCost > maxCalls = pure (bestOf seedCand frontier) - | length cands >= maxCands = pure (bestOf seedCand frontier) - | otherwise = case sampleParent seed (paretoFrontier frontier) of - Nothing -> pure (bestOf seedCand frontier) - Just (parent, seed') -> do - let parentProg = rebuild parent - idx = step `mod` nNodes - (fblog, _) <- captureFeedback train fbMetric parentProg - case drop idx paths of - (path : _) - | null (feedbackFor path fblog) -> - -- nothing to reflect on at this node; advance (capture cost only) - loop (step + 1) (calls + n) cands seed' frontier - _ -> do - child <- mutateNode proposer progSummary dataSummary fields fblog paths idx parentProg - rpt <- evaluatePure train metric child - let childCand = Candidate (foldParams child) (perEx rpt) (aggregateScore rpt) - frontier' = paretoFrontier (childCand : frontier) - loop (step + 1) (calls + stepCost) (childCand : cands) seed' frontier' + loop step cands seed frontier + | step >= stepCap = pure (bestOf seedCand frontier) + | length cands >= maxCands = pure (bestOf seedCand frontier) + | otherwise = do + fitsStep <- tryCharge meter stepCost + if not fitsStep + then pure (bestOf seedCand frontier) + else case sampleParent seed (paretoFrontier frontier) of + Nothing -> pure (bestOf seedCand frontier) + Just (parent, seed') -> do + let parentProg = rebuild parent + idx = step `mod` nNodes + (fblog, _) <- captureFeedback train fbMetric parentProg + case drop idx paths of + (path : _) + | null (feedbackFor path fblog) -> + -- nothing to reflect on at this node; the reserved + -- full-step budget is a conservative upper bound. + loop (step + 1) cands seed' frontier + _ -> do + child <- mutateNode proposer progSummary dataSummary fields fblog paths idx parentProg + rpt <- evaluatePure train metric child + let childCand = Candidate (foldParams child) (perEx rpt) (aggregateScore rpt) + frontier' = paretoFrontier (childCand : frontier) + loop (step + 1) (childCand : cands) seed' frontier' - best <- loop 0 n [seedCand] 1 [seedCand] - pure (freezeProgram (rebuild best)) + best <- loop 0 [seedCand] 1 [seedCand] + pure (freezeProgram (rebuild best)) -- | The per-example score vector from a report, in dataset order. perEx :: Report -> [Double] @@ -259,9 +260,7 @@ -- | The frontier candidate with the highest aggregate (earliest on ties); falls back -- to the seed if the frontier is somehow empty. bestOf :: Candidate -> [Candidate] -> Candidate -bestOf seedCand cs = case cs of - [] -> seedCand - (x : xs) -> foldl' (\b c -> if aggregate c > aggregate b then c else b) x xs +bestOf seedCand = foldl' (\b c -> if aggregate c > aggregate b then c else b) seedCand -- | A minimal program summary (EP-19's program describer is the richer source). fallbackProgramSummary :: Int -> Text diff --git a/src/Shikumi/Optimize/Instruction.hs b/src/Shikumi/Optimize/Instruction.hs --- a/src/Shikumi/Optimize/Instruction.hs +++ b/src/Shikumi/Optimize/Instruction.hs @@ -9,16 +9,18 @@ -- The proposer and its signal-gatherers are themselves ordinary shikumi 'Program's, so -- they are typed, cached, traced, and testable with the same stub-LM machinery as -- everything else — the optimizer is written in the framework it optimizes. The --- /current/ instruction is always retained as a candidate (the proposer guarantees it), --- so a node can never end up worse than where it started. +-- /current/ effective instruction (override, or signature base when no override is +-- present) is always retained as a candidate (the proposer guarantees it), and keeping +-- that candidate writes no redundant override, so a node can never end up worse than +-- where it started. -- --- __Budget.__ The grounded proposer makes @4 + proposalsPerNode@ LM calls per node --- (dataset summary, program describe, module describe, and one generation per proposal); --- scoring one candidate costs one LM call per dataset example. The search threads a --- running raw-call count and stops — returning the best found /so far/ — before either --- bound in the 'Budget' would be exceeded, so the recorded LM-call count never exceeds --- @maxLmCalls@. When the remaining budget cannot cover a node's full proposal, that node --- keeps its current instruction (no proposer call) rather than partially proposing. +-- __Budget.__ The grounded proposer reserves @4 + proposalsPerNode@ predicted LM +-- completions per node (dataset summary, program describe, module describe, and one +-- generation per proposal); scoring one candidate reserves one completion per dataset +-- example per predict node. The search stops — returning the best found /so far/ — +-- before either bound in the 'Budget' would be exceeded. When the remaining budget +-- cannot cover a node's full proposal, that node keeps its current instruction (no +-- proposer call) rather than partially proposing. -- -- This module re-points V1's blind proposer at EP-19's grounded surface; the old -- @ProposeIn@/@ProposeOut@/@proposeInstruction@ predictor is removed (the grounded @@ -31,14 +33,12 @@ import Control.Monad (foldM) import Data.Aeson (ToJSON) -import Data.Maybe (fromMaybe) -import Shikumi.Eval (datasetSize) import Shikumi.Optimize.Propose ( ProposeRequest (..), ProposeResult (..), proposeInstructions, ) -import Shikumi.Optimize.Search (freezeProgram, instructionAt, scoreOn, setNodeInstr) +import Shikumi.Optimize.Search (effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..)) import Shikumi.Program (foldParams) @@ -46,29 +46,31 @@ -- an explicit LM-call budget, using the grounded proposer to generate candidates. instructionSearch :: (ToJSON i, ToJSON o) => Int -> Budget -> Optimizer i o instructionSearch proposalsPerNode budget = Optimizer $ \train metric student -> do - let dsSize = datasetSize train - nNodes = length (foldParams student) + meter <- newBudgetMeter budget + let nNodes = length (foldParams student) -- The grounded proposer's fixed per-node LM-call cost (see module header). proposerCost = 4 + max 0 proposalsPerNode - -- Score candidate instructions for node @idx@, threading the call count and + -- Score candidate instructions for node @idx@, using the shared meter and -- the best-so-far; stop before the next scoring would exceed the budget. - scoreCands calls best idx prog cs = case cs of - [] -> pure (best, calls) - (c : rest) - | calls + dsSize > maxLmCalls budget -> pure (best, calls) - | otherwise -> do - s <- scoreOn train metric (setNodeInstr idx c prog) + scoreCands best idx prog cs = case cs of + [] -> pure best + (c : rest) -> do + ms <- meteredScore meter train metric (setNodeInstrIfNew idx c prog) + case ms of + Nothing -> pure best + Just s -> do let best' = case best of Nothing -> Just (c, s) Just (_, bs) -> if s > bs then Just (c, s) else best - scoreCands (calls + dsSize) best' idx prog rest + scoreCands best' idx prog rest -- Optimize one node, holding the others fixed. - stepNode (prog, calls) idx = do - let curInstr = fromMaybe "" (instructionAt idx prog) - (cands, calls1) <- - if calls + proposerCost <= maxLmCalls budget + stepNode prog idx = do + let curInstr = effectiveInstructionAt idx prog + fitsProposer <- tryCharge meter proposerCost + cands <- + if fitsProposer then do ProposeResult cs <- proposeInstructions @@ -83,10 +85,10 @@ tipIndex = 0, viewBatch = 2 } - pure (cs, calls + proposerCost) - else pure ([curInstr], calls) - (best, calls2) <- scoreCands calls1 Nothing idx prog cands - pure (setNodeInstr idx (maybe curInstr fst best) prog, calls2) + pure cs + else pure [curInstr] + best <- scoreCands Nothing idx prog cands + pure (setNodeInstrIfNew idx (maybe curInstr fst best) prog) - (final, _) <- foldM stepNode (student, 0) [0 .. nNodes - 1] + final <- foldM stepNode student [0 .. nNodes - 1] pure (freezeProgram final) diff --git a/src/Shikumi/Optimize/KNN.hs b/src/Shikumi/Optimize/KNN.hs --- a/src/Shikumi/Optimize/KNN.hs +++ b/src/Shikumi/Optimize/KNN.hs @@ -13,7 +13,8 @@ -- EP-15's pure @runEmbedding@ argument), not the @Embedding@ effect: an 'Embed' -- body's row is fixed to @(LLM, Error ShikumiError)@, so it cannot call @embedText@; -- all embedding-effect work happens at the caller, outside the node. The run-time --- form carries no @Params@ (it serializes as the empty vector, exactly like @react@). +-- form carries no @Params@ (it serializes as an @Embed@ shape with an empty vector, +-- exactly like @react@). module Shikumi.Optimize.KNN ( -- * Run-time form (the faithful analog) knnFewShot, @@ -100,7 +101,13 @@ in embed $ \i -> runProgram (withDemos (nearestDemos embedder k exs (toPrompt i)) student) i -- | Run-time KNN as an 'Optimizer': selection is by embedding geometry, not by --- score, so it consults neither the metric nor the LM at optimize time. +-- score, so it consults neither the metric nor the LM at optimize time and spends +-- zero optimizer LM calls. The result is a structure-changing @Embed@ wrapper +-- around the student. Its run-time selector closure is not persisted by +-- 'Shikumi.Compile.Serialize.encodeCompiled', so the saved state must be decoded +-- onto the same @knnDemos@ template; decode onto the plain student fails with a +-- shape mismatch. Use 'knnFewShotCentroid' when the desired artifact is only +-- serializable node parameters. knnFewShot :: (ToJSON i, ToJSON o, ToPrompt i) => (Text -> Vector Double) -> @@ -111,7 +118,7 @@ -- | Compile-time fallback: bake the @k@ training examples nearest the training -- centroid as a fixed demo set on every node, and freeze. A plain @Params@ artifact, --- no run-time embedder needed. +-- no run-time embedder needed, and zero optimizer LM calls are spent. knnFewShotCentroid :: (ToJSON i, ToJSON o, ToPrompt i) => (Text -> Vector Double) -> diff --git a/src/Shikumi/Optimize/LabeledFewShot.hs b/src/Shikumi/Optimize/LabeledFewShot.hs --- a/src/Shikumi/Optimize/LabeledFewShot.hs +++ b/src/Shikumi/Optimize/LabeledFewShot.hs @@ -5,10 +5,12 @@ -- -- No LM calls are made beyond scoring. Candidate sets are enumerated -- deterministically (all size-@k@ combinations of the training demos, in a fixed --- order, bounded by the budget) rather than randomly sampled, so the result is --- reproducible run to run — the tests rely on this. +-- order) rather than randomly sampled. 'labeledFewShotWith' reserves one scoring +-- cost per candidate before evaluation and stops at the shared 'Budget'; the +-- default 'labeledFewShot' uses 'defaultBudget'. module Shikumi.Optimize.LabeledFewShot ( labeledFewShot, + labeledFewShotWith, labeledCandidateSets, withDemos, ) @@ -18,16 +20,21 @@ import Data.Aeson (ToJSON, toJSON) import Data.Generics.Labels () import Shikumi.Eval (Dataset, Example (..), datasetExamples) -import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest) -import Shikumi.Optimize.Types (Optimizer (..), Scored (..), defaultBudget) +import Shikumi.Optimize.Search (freezeProgram, meteredScore, newBudgetMeter, selectBestMetered) +import Shikumi.Optimize.Types (Budget, Optimizer (..), Scored (..), defaultBudget) import Shikumi.Program (Demo (..), Program, mapParams) -- | Select the best size-@k@ set of labelled demonstrations from the training set. -- Multi-node programs receive the same demo set at every node (the DSPy default). labeledFewShot :: (ToJSON i, ToJSON o) => Int -> Optimizer i o -labeledFewShot k = Optimizer $ \train metric prog -> do +labeledFewShot = labeledFewShotWith defaultBudget + +-- | Select the best size-@k@ labelled demo set under an explicit budget. +labeledFewShotWith :: (ToJSON i, ToJSON o) => Budget -> Int -> Optimizer i o +labeledFewShotWith budget k = Optimizer $ \train metric prog -> do + meter <- newBudgetMeter budget let sets = labeledCandidateSets k train - best <- selectBest defaultBudget (\ds -> scoreOn train metric (withDemos ds prog)) sets + best <- selectBestMetered meter (\ds -> meteredScore meter train metric (withDemos ds prog)) sets pure $ case best of Nothing -> freezeProgram prog Just sc -> freezeProgram (withDemos (candidate sc) prog) diff --git a/src/Shikumi/Optimize/MIPRO.hs b/src/Shikumi/Optimize/MIPRO.hs --- a/src/Shikumi/Optimize/MIPRO.hs +++ b/src/Shikumi/Optimize/MIPRO.hs @@ -5,8 +5,9 @@ -- one node at a time, MIPROv2 searches /both/ axes — instruction × demo set — across -- all nodes jointly, in three phases: -- --- 1. __bootstrap__ candidate demo sets (per node) from the teacher's passing runs --- plus the labelled training pairs ('bootstrapDemoCandidates'); +-- 1. __bootstrap__ candidate demo sets (per node), retaining the node's current +-- demos at index 0, then adding the empty set, teacher passing runs, and the +-- labelled training pairs ('bootstrapDemoCandidates'); -- 2. __propose__ candidate instructions (per node) through EP-19's grounded -- proposer ('proposeInstructionCandidates'); and -- 3. __search__ the joint @(instruction × demoset)@ grid with cheap minibatch @@ -44,7 +45,6 @@ import Control.Monad (forM) import Data.Aeson (ToJSON) import Data.Aeson.Text (encodeToLazyText) -import Data.Maybe (fromMaybe) import Data.Text (Text) import Data.Text.Lazy qualified as TL import Effectful (Eff, (:>)) @@ -71,7 +71,7 @@ ProposeResult (..), proposeInstructions, ) -import Shikumi.Optimize.Search (freezeProgram, scoreOn) +import Shikumi.Optimize.Search (BudgetMeter, effectiveInstructionAt, freezeProgram, meteredScore, newBudgetMeter, setNodeInstrIfNew, tryCharge) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), defaultBudget) import Shikumi.Program ( Demo (..), @@ -95,7 +95,7 @@ data Miprov2Config = Miprov2Config { -- | instruction candidates proposed per node (incl. the current at index 0) numInstructCandidates :: !Int, - -- | demo-set candidates per node (incl. the empty set at index 0) + -- | demo-set candidates per node (incl. the node's current demos at index 0) numDemoCandidates :: !Int, -- | minibatch trials the search performs numTrials :: !Int, @@ -109,7 +109,7 @@ maxBootstrappedDemos :: !Int, -- | min metric score for a teacher run to contribute demos bootstrapThreshold :: !Double, - -- | hard LM-call / candidate ceiling (V1's 'Budget') + -- | hard predicted LM-completion / candidate ceiling (V1's 'Budget') budget :: !Budget } deriving stock (Eq, Show, Generic) @@ -134,9 +134,10 @@ -- | MIPROv2 with an explicit config and an explicit teacher program. miprov2With :: (ToJSON i, ToJSON o) => Miprov2Config -> Program i o -> Optimizer i o miprov2With cfg teacher = Optimizer $ \train metric student -> do - demoCands <- bootstrapDemoCandidates cfg teacher train metric student - instrCands <- proposeInstructionCandidates cfg student train demoCands - best <- searchJoint cfg train metric student instrCands demoCands + meter <- newBudgetMeter (budget cfg) + demoCands <- bootstrapDemoCandidatesWith meter cfg teacher train metric student + instrCands <- proposeInstructionCandidatesWith meter cfg student train demoCands + best <- searchJointWith meter cfg train metric student instrCands demoCands pure (freezeProgram best) -- --------------------------------------------------------------------------- @@ -144,14 +145,12 @@ -- --------------------------------------------------------------------------- -- | For each node (in @foldParams@ order), a list of candidate demo sets. Candidate --- 0 is always the empty set (so "no demos" is always reachable). The remaining sets --- come from the teacher's metric-passing runs (bootstrapped) and from the labelled --- training pairs (labelled demos, DSPy's @max_labeled_demos@) — recovered at the --- program-I/O level and attached to every node (the documented degradation while --- per-node trace recovery via EP-16 is deferred; for single-node programs the two --- coincide). +-- 0 is the node's current demos, making the baseline vector an identity. The empty +-- set remains reachable after that, followed by teacher metric-passing runs +-- (bootstrapped) and labelled training pairs (DSPy's @max_labeled_demos@) recovered at +-- the program-I/O level and attached to every node. bootstrapDemoCandidates :: - (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => Miprov2Config -> -- | teacher Program i o -> @@ -161,7 +160,23 @@ Program i o -> Eff es [[[Demo]]] bootstrapDemoCandidates cfg teacher train metric student = do + meter <- newBudgetMeter (budget cfg) + bootstrapDemoCandidatesWith meter cfg teacher train metric student + +bootstrapDemoCandidatesWith :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => + BudgetMeter -> + Miprov2Config -> + -- | teacher + Program i o -> + Dataset i o -> + Metric o -> + -- | student (for the node count) + Program i o -> + Eff es [[[Demo]]] +bootstrapDemoCandidatesWith meter cfg teacher train metric student = do let exs = datasetExamples train + teacherCost = max 1 (length (foldParams teacher)) keepIfPassing (Example inp expd) = ( do out <- runProgram teacher inp @@ -169,13 +184,23 @@ pure [recoverDemo inp out | s >= bootstrapThreshold cfg] ) `catchError` \_ (_ :: ShikumiError) -> pure [] - bootstrapped <- concat <$> mapM keepIfPassing exs + collect [] = pure [] + collect (ex : rest) = do + fits <- tryCharge meter teacherCost + if not fits + then pure [] + else do + kept <- keepIfPassing ex + (kept ++) <$> collect rest + bootstrapped <- collect exs let cap = max 1 (maxBootstrappedDemos cfg) labeledSet = take cap (map (\(Example i o) -> recoverDemo i o) exs) bootSet = take cap bootstrapped - sets = take (max 1 (numDemoCandidates cfg)) ([] : filter (not . null) [labeledSet, bootSet]) - nNodes = length (foldParams student) - pure (replicate nNodes sets) + nodeSets ps = + take + (max 1 (numDemoCandidates cfg)) + (dedup (demos ps : [] : filter (not . null) [labeledSet, bootSet])) + pure (map nodeSets (foldParams student)) -- --------------------------------------------------------------------------- -- Phase 2 — propose instruction candidates @@ -186,30 +211,48 @@ -- Uses EP-19's 'proposeInstructions'; the per-candidate tip starts at the /creative/ -- tip (@tipIndex = 1@) so a node always gets a strong proposal even at small N. proposeInstructionCandidates :: - (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es) => + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => Miprov2Config -> Program i o -> Dataset i o -> [[[Demo]]] -> Eff es [[Text]] -proposeInstructionCandidates cfg student train demoCands = - forM (zip3 [0 ..] (foldParams student) demoCands) $ \(k, ps, nodeDemoSets) -> do - let curInstr = fromMaybe "" (instructionOverride ps) +proposeInstructionCandidates cfg student train demoCands = do + meter <- newBudgetMeter (budget cfg) + proposeInstructionCandidatesWith meter cfg student train demoCands + +proposeInstructionCandidatesWith :: + (ToJSON i, ToJSON o, LLM :> es, Error ShikumiError :> es, Prim :> es) => + BudgetMeter -> + Miprov2Config -> + Program i o -> + Dataset i o -> + [[[Demo]]] -> + Eff es [[Text]] +proposeInstructionCandidatesWith meter cfg student train demoCands = + forM (zip [0 ..] demoCands) $ \(k, nodeDemoSets) -> do + let curInstr = effectiveInstructionAt k student demoTexts = map renderDemo (concat (take 1 (filter (not . null) nodeDemoSets))) - ProposeResult cs <- - proposeInstructions - train - ProposeRequest - { program = student, - targetNode = k, - currentInstruction = curInstr, - history = [], - bootstrappedDemos = demoTexts, - numCandidates = max 0 (numInstructCandidates cfg - 1), - tipIndex = 1, - viewBatch = 2 - } - pure cs + numProposals = max 0 (numInstructCandidates cfg - 1) + proposerCost = 4 + numProposals + fits <- tryCharge meter proposerCost + if not fits + then pure [curInstr] + else do + ProposeResult cs <- + proposeInstructions + train + ProposeRequest + { program = student, + targetNode = k, + currentInstruction = curInstr, + history = [], + bootstrappedDemos = demoTexts, + numCandidates = numProposals, + tipIndex = 1, + viewBatch = 2 + } + pure cs -- | Render a recovered 'Demo' as @ => @ for the proposal -- prompt's demo signal. @@ -227,9 +270,11 @@ -- | Search the joint per-node @(instruction × demoset)@ grid by greedy coordinate -- descent with minibatch screening, returning the best program found within the --- 'Budget'. Each trial screens the one-coordinate neighbours of the running best on a --- seeded minibatch, then full-evaluates the best-screened neighbour and accepts it --- only if it strictly improves the full score. +-- 'Budget'. Bootstrap teacher runs, grounded proposer calls, minibatch scoring, and +-- full scoring all reserve predicted cost against one meter in 'miprov2With'. Each +-- trial screens the one-coordinate neighbours of the running best on a seeded +-- minibatch, then full-evaluates the best-screened neighbour and accepts it only if +-- it strictly improves the full score. searchJoint :: (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => Miprov2Config -> @@ -242,53 +287,73 @@ [[[Demo]]] -> Eff es (Program i o) searchJoint cfg train metric student instrCands demoCands = do + meter <- newBudgetMeter (budget cfg) + searchJointWith meter cfg train metric student instrCands demoCands + +-- | Meter-aware implementation of 'searchJoint'. Later MIPRO phases share a meter +-- with this helper so all three phases consume one budget. +searchJointWith :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + BudgetMeter -> + Miprov2Config -> + Dataset i o -> + Metric o -> + Program i o -> + -- | per-node instruction candidates + [[Text]] -> + -- | per-node demo-set candidates + [[[Demo]]] -> + Eff es (Program i o) +searchJointWith meter cfg train metric student instrCands demoCands = do let nNodes = length (foldParams student) dsSize = datasetSize train mbSize = max 1 (min dsSize (minibatchSize cfg)) - maxCalls = maxLmCalls (budget cfg) apply = applyVec instrCands demoCands student baseVec = replicate nNodes (0, 0) - if dsSize > maxCalls - then pure (apply baseVec) - else do - baseFull <- scoreOn train metric (apply baseVec) - let screen _t calls acc [] = pure (reverse acc, calls) - screen t calls acc (v : vs) - | calls + mbSize > maxCalls = pure (reverse acc, calls) - | otherwise = do - s <- scoreOn (minibatchAt mbSize (t * 7 + length acc) train) metric (apply v) - screen t (calls + mbSize) ((v, s) : acc) vs - go t calls bestVec bestFull + mBaseFull <- meteredScore meter train metric (apply baseVec) + case mBaseFull of + Nothing -> pure (apply baseVec) + Just baseFull -> do + let screen _t acc [] = pure (reverse acc) + screen t acc (v : vs) = do + ms <- meteredScore meter (minibatchAt mbSize (t * 7 + length acc) train) metric (apply v) + case ms of + Nothing -> pure (reverse acc) + Just s -> screen t ((v, s) : acc) vs + go t bestVec bestFull | t >= numTrials cfg = pure bestVec | otherwise = do - (scored, calls1) <- screen t calls [] (oneCoordNeighbors instrCands demoCands bestVec) + scored <- screen t [] (oneCoordNeighbors instrCands demoCands bestVec) case bestByScore scored of Nothing -> pure bestVec - Just propVec - | calls1 + dsSize > maxCalls -> pure bestVec - | otherwise -> do - fs <- scoreOn train metric (apply propVec) + Just propVec -> do + mFull <- meteredScore meter train metric (apply propVec) + case mFull of + Nothing -> pure bestVec + Just fs -> if fs > bestFull - then go (t + 1) (calls1 + dsSize) propVec fs - else go (t + 1) (calls1 + dsSize) bestVec bestFull - best <- go 0 dsSize baseVec baseFull + then go (t + 1) propVec fs + else go (t + 1) bestVec bestFull + best <- go 0 baseVec baseFull pure (apply best) -- | Apply a joint vector to the student: set each node's instruction and demos to the --- selected candidates. +-- selected candidates. Index 0 on either axis is an identity by construction. applyVec :: [[Text]] -> [[[Demo]]] -> Program i o -> JointVec -> Program i o applyVec instrCands demoCands prog vec = foldl step prog (zip [0 ..] vec) where step p (k, (i, d)) = - mapParamsAt - k - ( \ps -> - ps - { instructionOverride = Just (instrCands !! k !! i), - demos = demoCands !! k !! d - } - ) - p + setNodeInstrIfNew k (instrCands !! k !! i) $ + mapParamsAt k (\ps -> ps {demos = demoCands !! k !! d}) p + +-- | Order-preserving de-duplication. +dedup :: (Eq a) => [a] -> [a] +dedup = go [] + where + go _ [] = [] + go seen (x : xs) + | x `elem` seen = go seen xs + | otherwise = x : go (x : seen) xs -- | The one-coordinate neighbours of a vector: each move changes exactly one node's -- instruction or demo index to a non-current alternative. diff --git a/src/Shikumi/Optimize/RandomSearch.hs b/src/Shikumi/Optimize/RandomSearch.hs --- a/src/Shikumi/Optimize/RandomSearch.hs +++ b/src/Shikumi/Optimize/RandomSearch.hs @@ -20,10 +20,10 @@ import Data.Generics.Labels () import Data.List (sortBy) import Data.Ord (comparing) -import Shikumi.Compile.Types (compiledProgram) import Shikumi.Eval (Dataset, dataset, datasetExamples) -import Shikumi.Optimize.Bootstrap (bootstrapFewShotWith, defaultBootstrapConfig) -import Shikumi.Optimize.Search (freezeProgram, scoreOn, selectBest) +import Shikumi.Optimize.Bootstrap (bootstrapKeptDemos, defaultBootstrapConfig) +import Shikumi.Optimize.LabeledFewShot (withDemos) +import Shikumi.Optimize.Search (freezeProgram, meteredScore, newBudgetMeter, selectBestMetered) import Shikumi.Optimize.Types (Budget (..), Optimizer (..), Scored (..)) import Shikumi.Program (Program) @@ -66,7 +66,10 @@ (x : _) -> lo + (x `mod` span') [] -> lo --- | 'bootstrapRandomSearch' with explicit tunables. +-- | 'bootstrapRandomSearch' with explicit tunables. One 'Budget' covers all seed +-- bootstrap teacher runs and the final candidate scoring pass; when the meter is +-- exhausted, later seeds or scoring candidates are skipped and the best scored +-- candidate so far is returned. bootstrapRandomSearchWith :: (ToJSON i, ToJSON o) => RandomSearchConfig -> @@ -77,14 +80,15 @@ Budget -> Optimizer i o bootstrapRandomSearchWith cfg teacher numCandidates budget = Optimizer $ \train metric student -> do + meter <- newBudgetMeter budget let seeds = [1 .. max 1 numCandidates] candidateFor seed = do let cfg' = defaultBootstrapConfig & #maxBootstrappedDemos .~ sizeFor cfg seed - opt = bootstrapFewShotWith cfg' teacher budget - compiledProgram <$> runOptimizer opt (shuffle seed train) metric student + demos <- bootstrapKeptDemos cfg' meter teacher (shuffle seed train) metric + pure (withDemos demos student) seeded <- mapM candidateFor seeds let cands = student : seeded -- zero-shot baseline first - best <- selectBest budget (scoreOn train metric) cands + best <- selectBestMetered meter (\p -> meteredScore meter train metric p) cands pure $ case best of Nothing -> freezeProgram student Just sc -> freezeProgram (candidate sc) diff --git a/src/Shikumi/Optimize/Search.hs b/src/Shikumi/Optimize/Search.hs --- a/src/Shikumi/Optimize/Search.hs +++ b/src/Shikumi/Optimize/Search.hs @@ -10,26 +10,37 @@ module Shikumi.Optimize.Search ( selectBest, scoreOn, + BudgetMeter (..), + newBudgetMeter, + tryCharge, + scoringCost, + meteredScore, + selectBestMetered, + withLmCallCount, freezeProgram, setNodeInstr, + setNodeInstrIfNew, instructionAt, + effectiveInstructionAt, ) where import Control.Lens ((&), (?~)) import Data.Generics.Labels () +import Data.Maybe (fromMaybe) import Data.Text (Text) import Effectful (Eff, (:>)) import Effectful.Concurrent (Concurrent) +import Effectful.Dispatch.Dynamic (interpose) import Effectful.Error.Static (Error) -import Effectful.Prim (Prim) +import Effectful.Prim.IORef (IORef, Prim, atomicModifyIORef', newIORef, readIORef) import Shikumi.Compile.Types (CompiledProgram (..)) import Shikumi.Effect.Time (Time) import Shikumi.Error (ShikumiError) -import Shikumi.Eval (Dataset, Metric, Report (aggregateScore), evaluatePure) -import Shikumi.LLM (LLM) +import Shikumi.Eval (Dataset, Metric, Report (aggregateScore), datasetSize, evaluatePure) +import Shikumi.LLM (LLM (..), complete, stream) import Shikumi.Optimize.Types (Budget (..), Scored (..)) -import Shikumi.Program (Params (..), Program, foldParams, mapParamsAt) +import Shikumi.Program (Params (..), Program, foldParams, mapParamsAt, nodeInstructionsIndexed) -- | Score every candidate (left to right), stopping once the candidate budget is -- hit, and return the best by score (ties: earliest wins). The scorer is the @@ -54,6 +65,112 @@ bestOf [] = Nothing bestOf (x : xs) = Just (foldl' (\b c -> if score c > score b then c else b) x xs) +-- | Mutable spend-tracking for one optimizer run. The call counter tracks +-- predicted LM completions already reserved; the candidate counter tracks +-- scored candidates. +data BudgetMeter = BudgetMeter + { meterBudget :: !Budget, + meterCalls :: !(IORef Int), + meterCands :: !(IORef Int) + } + +-- | Create a fresh meter for one optimizer invocation. +newBudgetMeter :: (Prim :> es) => Budget -> Eff es BudgetMeter +newBudgetMeter budget = do + calls <- newIORef 0 + cands <- newIORef 0 + pure BudgetMeter {meterBudget = budget, meterCalls = calls, meterCands = cands} + +-- | Charge @n@ predicted LM calls only when doing so keeps the run inside +-- 'maxLmCalls'. Non-positive charges are free and succeed. +tryCharge :: (Prim :> es) => BudgetMeter -> Int -> Eff es Bool +tryCharge _ n | n <= 0 = pure True +tryCharge meter n = + atomicModifyIORef' (meterCalls meter) $ \calls -> + if calls + n <= maxLmCalls (meterBudget meter) + then (calls + n, True) + else (calls, False) + +-- | The predicted LM-call cost of scoring @p@ over @ds@ once: one call per +-- example per 'Predict' node. Wrappers that re-run the LM can spend more. +scoringCost :: Dataset i o -> Program i o -> Int +scoringCost ds p = datasetSize ds * max 1 (length (foldParams p)) + +-- | Score one candidate program under a meter. Returns 'Nothing' without +-- spending when either the candidate ceiling is reached or the predicted scoring +-- cost no longer fits. +meteredScore :: + (LLM :> es, Concurrent :> es, Error ShikumiError :> es, Time :> es, Prim :> es) => + BudgetMeter -> + Dataset i o -> + Metric o -> + Program i o -> + Eff es (Maybe Double) +meteredScore meter ds metric prog = do + fitsCandidate <- tryBumpCandidate meter + if not fitsCandidate + then pure Nothing + else do + fitsCalls <- tryCharge meter (scoringCost ds prog) + if fitsCalls + then Just <$> scoreOn ds metric prog + else do + unbumpCandidate meter + pure Nothing + +-- | Fold candidates left-to-right under a scorer that can stop the search by +-- returning 'Nothing'. Ties keep the earliest candidate. +selectBestMetered :: + BudgetMeter -> + (cand -> Eff es (Maybe Double)) -> + [cand] -> + Eff es (Maybe (Scored cand)) +selectBestMetered _ scorer = go Nothing + where + go best [] = pure best + go best (cand : rest) = do + ms <- scorer cand + case ms of + Nothing -> pure best + Just s -> + let scored = Scored cand s + best' = case best of + Nothing -> Just scored + Just old -> Just (if score scored > score old then scored else old) + in go best' rest + +-- | Run an action while counting every LLM operation exactly. This is used where +-- the spender is opaque to the budget predictor. +withLmCallCount :: (LLM :> es, Prim :> es) => Eff es a -> Eff es (a, Int) +withLmCallCount act = do + ref <- newIORef 0 + result <- + interpose + ( \_ -> \case + Complete m c o -> do + bump ref + complete m c o + Stream m c o -> do + bump ref + stream m c o + ) + act + count <- readIORef ref + pure (result, count) + where + bump ref = atomicModifyIORef' ref (\n -> (n + 1, ())) + +tryBumpCandidate :: (Prim :> es) => BudgetMeter -> Eff es Bool +tryBumpCandidate meter = + atomicModifyIORef' (meterCands meter) $ \cands -> + if cands < max 0 (maxCandidates (meterBudget meter)) + then (cands + 1, True) + else (cands, False) + +unbumpCandidate :: (Prim :> es) => BudgetMeter -> Eff es () +unbumpCandidate meter = + atomicModifyIORef' (meterCands meter) $ \cands -> (max 0 (cands - 1), ()) + -- | Score a candidate program against a dataset and pure metric: run EP-8's -- @evaluate@ and take the aggregate. One call scores one program over the whole -- dataset (one LM call per example). @@ -76,6 +193,24 @@ -- @instructionSearch@ and @copro@. setNodeInstr :: Int -> Text -> Program i o -> Program i o setNodeInstr idx instr = mapParamsAt idx (\ps -> ps & #instructionOverride ?~ instr) + +-- | The instruction node @idx@ actually runs with: its 'instructionOverride' when +-- set, otherwise its signature's base instruction. This mirrors the runtime +-- precedence in 'Shikumi.Program.effectiveSignature'. Out-of-range indices yield +-- the empty string. +effectiveInstructionAt :: Int -> Program i o -> Text +effectiveInstructionAt idx prog = + case drop idx (zip (nodeInstructionsIndexed prog) (foldParams prog)) of + ((base, ps) : _) -> fromMaybe base (instructionOverride ps) + [] -> "" + +-- | Set node @idx@'s instruction override unless @instr@ is already its effective +-- instruction. Keeping the current signature instruction therefore leaves +-- 'instructionOverride' as 'Nothing' instead of serializing a redundant override. +setNodeInstrIfNew :: Int -> Text -> Program i o -> Program i o +setNodeInstrIfNew idx instr prog + | instr == effectiveInstructionAt idx prog = prog + | otherwise = setNodeInstr idx instr prog -- | The instruction override stored at node @idx@ (in @foldParams@ order), if any. instructionAt :: Int -> Program i o -> Maybe Text diff --git a/src/Shikumi/Optimize/Types.hs b/src/Shikumi/Optimize/Types.hs --- a/src/Shikumi/Optimize/Types.hs +++ b/src/Shikumi/Optimize/Types.hs @@ -10,6 +10,15 @@ -- never changes a program's structure or types — only its parameters — so the -- optimized program is the same typed function, merely better-behaved. -- +-- Two optimizers are explicit structure-changing exceptions. 'Shikumi.Optimize.KNN.knnFewShot' +-- returns an @Embed@ wrapper that selects demos at run time from an opaque closure; +-- persist the underlying student or use 'Shikumi.Optimize.KNN.knnFewShotCentroid' +-- when a plain parameter artifact is required. 'Shikumi.Optimize.Ensemble.ensembleSearch' +-- returns an @Ensemble@ over member programs; its saved state can only be loaded +-- onto the matching ensemble template, because the reducer closure and member +-- structure are part of the program held in code rather than the saved parameter +-- state. +-- -- __Why a record of functions, not a typeclass.__ The four optimizers differ in -- what extra inputs they carry (a teacher program, a budget, an inner optimizer -- for ensembling) but share one driver signature. A @newtype@ whose field is the @@ -66,11 +75,21 @@ Eff es (CompiledProgram i o) } --- | A hard, explicit bound on search cost. The optimizers that issue LM calls --- count the calls they make (proposer calls + per-candidate scoring evaluations, --- each scoring evaluation costing one LM call per dataset example) and stop — --- returning the best candidate found /so far/ — before any bound is exceeded, so --- they never silently produce an unscored program and never blow a cost ceiling. +-- | A hard, explicit bound on optimizer search cost. The accounting unit is a +-- predicted LM completion: scoring one candidate costs one completion per dataset +-- example per 'Shikumi.Program.Predict' node; a teacher run costs one completion +-- per 'Shikumi.Program.Predict' node; and a grounded proposer call costs @4 + N@ +-- completions for @N@ requested proposals. Optimizers reserve this predicted cost +-- before spending it and return the best candidate found so far when the next +-- spend would exceed a ceiling. +-- +-- This prediction is exact for plain predict/compose programs. Wrappers that can +-- re-invoke the LM internally, such as retry, majority-vote, ensemble, or embed +-- bodies, make the prediction a lower bound. 'Shikumi.Optimize.Ensemble.ensembleSearchWith' +-- is the exception that counts member optimizers exactly, but enforces its ceiling +-- only between members, so it may overshoot by one member's cost. If a budget is +-- too small to score even the seed candidate, optimizers return the input program +-- unscored. data Budget = Budget { -- | ceiling on raw LM calls the optimizer may make (proposals + scoring) maxLmCalls :: !Int, diff --git a/test/AcceptanceSpec.hs b/test/AcceptanceSpec.hs --- a/test/AcceptanceSpec.hs +++ b/test/AcceptanceSpec.hs @@ -12,6 +12,9 @@ -- reproducible. module AcceptanceSpec (tests) where +import Control.Monad (when) +import Data.IORef (newIORef, readIORef) +import Data.Text qualified as T import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) @@ -20,13 +23,22 @@ import Shikumi.Compile.ZeroShot (zeroShot) import Shikumi.Effect.Time (Time, runTime) import Shikumi.Error (ShikumiError) -import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.Eval (Dataset, Prediction, Score, boolScore, dataset, exactMatch, example, predictionPrimary) import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Program) -import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import StubLM + ( Label (..), + Sentence (..), + ruleInstruction, + ruled, + runStubLM, + runStubLMCapturing, + sentimentPipeline, + sentimentProg, + ) import Test.Tasty (TestTree, testGroup) -import Test.Tasty.HUnit (assertBool, assertFailure, testCase) +import Test.Tasty.HUnit (Assertion, assertBool, assertFailure, testCase, (@?=)) runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act @@ -86,7 +98,30 @@ assertBool ("before should be low, got " <> show before) (before < 0.5) improved "bootstrapFewShot" before afterBoot improved "labeledFewShot" before afterLabeled - improved "instructionSearch" before afterInstr + improved "instructionSearch" before afterInstr, + testGroup + "seeding over multi-node programs" + [ testCase "instructionSearch optimizes node 0 without changing node 1" $ + checkPipeline "instructionSearch" (instructionSearch 3 defaultBudget) True, + testCase "copro optimizes node 0 without changing node 1" $ + checkPipeline "copro" (copro defaultCoproConfig) True, + testCase "miprov2 keeps the echo node's effective instruction" $ + checkPipeline "miprov2" (miprov2 Miprov2Light) False, + testCase "gepa keeps the echo node's effective instruction" $ + checkPipeline "gepa" (gepa reflectiveProposer fbMetric defaultBudget) False, + testCase "instructionSearch scores the true student as a candidate" $ do + ref <- newIORef [] + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCapturing ref (optimize (instructionSearch 1 defaultBudget) trainset exactMatch ruled) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + captured <- readIORef ref + assertBool + "expected a scored request rendered with the RULE signature instruction" + (any (\txt -> ruleInstruction `T.isInfixOf` txt && "text: good" `T.isInfixOf` txt) captured) + ] ] where -- Strictly higher than before, and at least a floor of 0.75, so the test fails @@ -98,3 +133,24 @@ assertBool (name <> ": expected " <> show after <> " >= 0.75") (after >= 0.75) + +fbMetric :: Label -> Prediction Label -> (Score, T.Text) +fbMetric expd p = + let correct = expd == predictionPrimary p + in (boolScore correct, if correct then "" else "be more specific") + +checkPipeline :: String -> Optimizer Sentence Label -> Bool -> Assertion +checkPipeline name opt requireFloor = do + res <- runStub $ do + before <- scoreOn heldout exactMatch sentimentPipeline + cp <- optimize opt trainset exactMatch sentimentPipeline + let out = compiledProgram cp + after <- scoreOn heldout exactMatch out + pure (before, after, effectiveInstructionAt 1 out) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after, node1Instruction) -> do + assertBool (name <> ": expected " <> show after <> " >= " <> show before) (after >= before) + when requireFloor $ + assertBool (name <> ": expected " <> show after <> " >= 0.75") (after >= 0.75) + node1Instruction @?= "Echo the sentiment label unchanged." diff --git a/test/BootstrapSpec.hs b/test/BootstrapSpec.hs --- a/test/BootstrapSpec.hs +++ b/test/BootstrapSpec.hs @@ -9,6 +9,7 @@ -- teacher correct on, and the mislabelled one is excluded. module BootstrapSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) @@ -22,9 +23,9 @@ import Shikumi.Optimize import Shikumi.Program (Demo (..), Params (..), Program, programParams) import Shikumi.Schema (fromModel) -import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentPipeline, sentimentProg) import Test.Tasty (TestTree, testGroup) -import Test.Tasty.HUnit (assertFailure, testCase, (@?=)) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act @@ -44,6 +45,15 @@ example (Sentence "good soup") (Label "negative") ] +budgetTrainset :: Dataset Sentence Label +budgetTrainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative"), + example (Sentence "good show") (Label "positive"), + example (Sentence "bad play") (Label "negative") + ] + -- | The demos baked into a compiled single-node program. compiledDemos :: CompiledProgram Sentence Label -> [Demo] compiledDemos cp = case programParams (compiledProgram cp) of @@ -68,5 +78,16 @@ outs = traverse (\d -> fromModel (output d)) ds :: Either ShikumiError [Label] length ds @?= 2 ins @?= Right [Sentence "good film", Sentence "bad film"] - outs @?= Right [Label "positive", Label "negative"] + outs @?= Right [Label "positive", Label "negative"], + testCase "bootstrapFewShot charges a teacher run by predict-node count" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 6, maxCandidates = 32} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (bootstrapFewShot sentimentPipeline budget) budgetTrainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 6, got " <> show count) (count > 0 && count <= 6) ] diff --git a/test/EnsembleSpec.hs b/test/EnsembleSpec.hs --- a/test/EnsembleSpec.hs +++ b/test/EnsembleSpec.hs @@ -10,6 +10,7 @@ module EnsembleSpec (tests) where import Data.Aeson (toJSON) +import Data.IORef (newIORef, readIORef) import Data.Text (Text) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) @@ -23,7 +24,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Demo (..), Program, ProgramShape (..), programShape) -import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import StubLM (Label (..), Sentence (..), runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -88,5 +89,18 @@ Left e -> assertFailure ("unexpected error: " <> show e) Right cp -> case programShape (compiledProgram cp) of ShapeEnsemble subs -> length subs @?= 3 - other -> assertFailure ("expected a 3-member ShapeEnsemble, got " <> show other) + other -> assertFailure ("expected a 3-member ShapeEnsemble, got " <> show other), + testCase "ensembleSearchWith enforces the budget between members" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 1, maxCandidates = 32} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (ensembleSearchWith budget 5 (labeledFewShot 1)) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right cp -> case programShape (compiledProgram cp) of + ShapeEnsemble subs -> length subs @?= 1 + other -> assertFailure ("expected a 1-member ShapeEnsemble, got " <> show other) + count <- readIORef ref + count @?= 4 ] diff --git a/test/GepaSpec.hs b/test/GepaSpec.hs --- a/test/GepaSpec.hs +++ b/test/GepaSpec.hs @@ -16,7 +16,8 @@ import Shikumi.Eval (Dataset, boolScore, dataset, exactMatch, example, predictionPrimary) import Shikumi.LLM (LLM) import Shikumi.Optimize - ( Candidate (..), + ( Budget (..), + Candidate (..), FeedbackMetric, captureFeedback, defaultBudget, @@ -28,6 +29,7 @@ reflectiveProposer, sampleParent, scoreOn, + withLmCallCount, ) import Shikumi.Program (Params (..), foldParams, nodeFieldsIndexed, programParams) import Shikumi.Trace.Feedback (feedbackFor) @@ -65,7 +67,7 @@ runGepa act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runGepaStubLM act tests :: TestTree -tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, roundTrips] +tests = testGroup "Gepa" [paretoPure, feedbackCapture, reflectiveMutation, heldoutLift, budgetGate, roundTrips] -- --------------------------------------------------------------------------- -- Pure Pareto-frontier tests @@ -153,6 +155,17 @@ assertBool ("before should be low, got " <> show before) (before < 0.5) assertBool ("after " <> show after <> " should exceed before " <> show before) (after > before) assertBool ("after should reach the floor, got " <> show after) (after >= 0.75) + +budgetGate :: TestTree +budgetGate = + testCase "gepa returns the student without seed evaluation when the budget is too small" $ do + let tiny = Budget {maxLmCalls = 1, maxCandidates = 4} + res <- runGepa (withLmCallCount (optimize (gepa reflectiveProposer fbMetric tiny) trainset exactMatch sentimentProg)) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (cp, calls) -> do + calls @?= 0 + programParams (compiledProgram cp) @?= programParams sentimentProg roundTrips :: TestTree roundTrips = diff --git a/test/InstructionSpec.hs b/test/InstructionSpec.hs --- a/test/InstructionSpec.hs +++ b/test/InstructionSpec.hs @@ -20,7 +20,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Params (..), programParams) -import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentProg) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentPipeline, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -57,5 +57,18 @@ count <- readIORef ref assertBool ("expected 0 < calls <= 6, got " <> show count) - (count > 0 && count <= 6) + (count > 0 && count <= 6), + testCase "charges scoring by dataset size times predict nodes" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 11, maxCandidates = 32} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (instructionSearch 1 budget) trainset exactMatch sentimentPipeline) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool + ("expected 0 < calls <= 11 for two-node scoring, got " <> show count) + (count > 0 && count <= 11) ] diff --git a/test/KNNSpec.hs b/test/KNNSpec.hs --- a/test/KNNSpec.hs +++ b/test/KNNSpec.hs @@ -23,7 +23,7 @@ import Shikumi.Error (ShikumiError) import Shikumi.Eval (Dataset, dataset, datasetExamples, exactMatch, example) import Shikumi.LLM (LLM) -import Shikumi.Optimize (knnDemos, knnFewShotCentroid, nearestDemos, optimize) +import Shikumi.Optimize (knnDemos, knnFewShot, knnFewShotCentroid, nearestDemos, optimize, withLmCallCount) import Shikumi.Optimize.Search (freezeProgram) import Shikumi.Program (Demo (..), Params (..), programParams, runProgram) import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) @@ -74,7 +74,7 @@ runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act tests :: TestTree -tests = testGroup "KNN" [inputNearest, centroidNearest, runtimeRuns, roundTrips] +tests = testGroup "KNN" [inputNearest, centroidNearest, runtimeRuns, optimizeBudget, roundTrips] -- --------------------------------------------------------------------------- -- M1 — selection @@ -105,6 +105,18 @@ res <- runStub (runProgram prog (Sentence "france capital quiz")) assertBool "the Embed-wrapped KNN program runs" (either (const False) (const True) res) +optimizeBudget :: TestTree +optimizeBudget = + testCase "KNN optimizers make zero LM calls at optimize time" $ do + res <- + runStub $ do + (_, runtimeCalls) <- withLmCallCount (optimize (knnFewShot stubEmbed 2) balanced exactMatch sentimentProg) + (_, centroidCalls) <- withLmCallCount (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg) + pure (runtimeCalls, centroidCalls) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right counts -> counts @?= (0, 0) + -- --------------------------------------------------------------------------- -- M3 — serialization round-trip -- --------------------------------------------------------------------------- @@ -119,6 +131,11 @@ case decodeCompiledOnto (knnDemos stubEmbed 2 balanced sentimentProg) (encodeCompiled compiled) of Left err -> assertFailure ("decode failed: " <> err) Right cp' -> programParams (compiledProgram cp') @?= [], + testCase "run-time KNN decode onto the plain student fails with a shape message" $ do + let compiled = freezeProgram (knnDemos stubEmbed 2 balanced sentimentProg) + case decodeCompiledOnto sentimentProg (encodeCompiled compiled) of + Left err -> assertBool ("shape mismatch mentioned: " <> err) ("shape mismatch" `T.isInfixOf` T.pack err) + Right _ -> assertFailure "expected decode onto the plain student to fail", testCase "centroid KNN baked demos round-trip onto the bare student" $ do res <- runStub (optimize (knnFewShotCentroid stubEmbed 2) skewed exactMatch sentimentProg) case res of diff --git a/test/LabeledFewShotSpec.hs b/test/LabeledFewShotSpec.hs --- a/test/LabeledFewShotSpec.hs +++ b/test/LabeledFewShotSpec.hs @@ -9,6 +9,7 @@ -- is a perfect 1.0. module LabeledFewShotSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) @@ -20,7 +21,7 @@ import Shikumi.LLM (LLM) import Shikumi.Optimize import Shikumi.Program (Demo, Params (..), programParams) -import StubLM (Label (..), Sentence (..), runStubLM, sentimentProg) +import StubLM (Label (..), Sentence (..), runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -70,5 +71,16 @@ (maximum candScores > minimum candScores) assertBool ("best candidate should be perfect, got " <> show chosenScore) - (chosenScore == 1.0) + (chosenScore == 1.0), + testCase "labeledFewShotWith respects maxLmCalls while scoring candidates" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 6, maxCandidates = 100} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (labeledFewShotWith budget 2) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 6, got " <> show count) (count > 0 && count <= 6) ] diff --git a/test/Main.hs b/test/Main.hs --- a/test/Main.hs +++ b/test/Main.hs @@ -17,6 +17,8 @@ import OptimizeSpec qualified import ProposeSpec qualified import RandomSearchSpec qualified +import SearchSpec qualified +import SeedingSpec qualified import Test.Tasty (defaultMain, testGroup) main :: IO () @@ -35,5 +37,7 @@ GepaSpec.tests, KNNSpec.tests, RandomSearchSpec.tests, + SearchSpec.tests, + SeedingSpec.tests, AcceptanceSpec.tests ] diff --git a/test/Miprov2Spec.hs b/test/Miprov2Spec.hs --- a/test/Miprov2Spec.hs +++ b/test/Miprov2Spec.hs @@ -28,6 +28,7 @@ bootstrapDemoCandidates, defaultBudget, instructionSearch, + miprov2, miprov2Auto, miprov2With, optimize, @@ -179,7 +180,28 @@ assertBool ("miprov2 " <> show afterMipro <> " should beat instructionSearch " <> show afterInstr) (afterMipro > afterInstr) - assertBool ("miprov2 should reach the floor, got " <> show afterMipro) (afterMipro >= 0.75) + assertBool ("miprov2 should reach the floor, got " <> show afterMipro) (afterMipro >= 0.75), + testCase "miprov2With shares one budget across bootstrap, proposal, and search" $ do + ref <- newIORef (0 :: Int) + let cfg = cfgLight & #budget .~ Budget {maxLmCalls = 4, maxCandidates = 32} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runJointStubLMCounting ref (optimize (miprov2With cfg sentimentProg) jointTrain exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 4, got " <> show count) (count > 0 && count <= 4), + testCase "miprov2 preset stays within the default budget" $ do + ref <- newIORef (0 :: Int) + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runJointStubLMCounting ref (optimize (miprov2 Miprov2Light) jointTrain exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 200, got " <> show count) (count > 0 && count <= maxLmCalls defaultBudget) ] serialize :: TestTree diff --git a/test/RandomSearchSpec.hs b/test/RandomSearchSpec.hs --- a/test/RandomSearchSpec.hs +++ b/test/RandomSearchSpec.hs @@ -5,6 +5,7 @@ -- the best compiled output round-trips through serialization. module RandomSearchSpec (tests) where +import Data.IORef (newIORef, readIORef) import Effectful (Eff, IOE, runEff) import Effectful.Concurrent (Concurrent, runConcurrent) import Effectful.Error.Static (Error, runErrorNoCallStack) @@ -16,7 +17,8 @@ import Shikumi.Eval (Dataset, dataset, exactMatch, example) import Shikumi.LLM (LLM) import Shikumi.Optimize - ( bootstrapFewShot, + ( Budget (..), + bootstrapFewShot, bootstrapRandomSearch, defaultBudget, optimize, @@ -24,7 +26,7 @@ setNodeInstr, ) import Shikumi.Program (Program) -import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, sentimentProg) +import StubLM (Label (..), Sentence (..), ruleInstruction, runStubLM, runStubLMCounting, sentimentProg) import Test.Tasty (TestTree, testGroup) import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) @@ -57,7 +59,7 @@ runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act tests :: TestTree -tests = testGroup "RandomSearch" [bestOfNAtLeastSingle, reproducible, roundTrips] +tests = testGroup "RandomSearch" [bestOfNAtLeastSingle, reproducible, budgetRespected, roundTrips] bestOfNAtLeastSingle :: TestTree bestOfNAtLeastSingle = @@ -89,6 +91,20 @@ case res of Left e -> assertFailure ("unexpected error: " <> show e) Right (sa, sb) -> sa @?= sb + +budgetRespected :: TestTree +budgetRespected = + testCase "bootstrapRandomSearch shares one budget across seeds and scoring" $ do + ref <- newIORef (0 :: Int) + let budget = Budget {maxLmCalls = 20, maxCandidates = 100} + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref (optimize (bootstrapRandomSearch teacher 5 budget) trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + count <- readIORef ref + assertBool ("expected 0 < calls <= 20, got " <> show count) (count > 0 && count <= 20) roundTrips :: TestTree roundTrips = diff --git a/test/SearchSpec.hs b/test/SearchSpec.hs new file mode 100644 --- /dev/null +++ b/test/SearchSpec.hs @@ -0,0 +1,114 @@ +{-# LANGUAGE TypeApplications #-} + +-- | Unit tests for the shared optimizer search helpers. +module SearchSpec (tests) where + +import Data.IORef (newIORef, readIORef) +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Effectful.Prim.IORef qualified as EIORef +import Shikumi.Combinator ((>>>)) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.LLM (LLM) +import Shikumi.Module (predict) +import Shikumi.Optimize.Search + ( BudgetMeter (..), + effectiveInstructionAt, + instructionAt, + meteredScore, + newBudgetMeter, + scoreOn, + selectBestMetered, + setNodeInstr, + setNodeInstrIfNew, + tryCharge, + withLmCallCount, + ) +import Shikumi.Optimize.Types (Budget (..), Scored (..)) +import Shikumi.Program (Program) +import Shikumi.Signature (Signature, mkSignature) +import StubLM (Label (..), Sentence (..), ruleInstruction, ruled, runStubLM, runStubLMCounting, sentimentProg) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "bad film") (Label "negative") + ] + +echoSig :: Signature Label Label +echoSig = mkSignature "Echo the sentiment label unchanged." + +twoNode :: Program Sentence Label +twoNode = sentimentProg >>> predict echoSig + +tests :: TestTree +tests = + testGroup + "Search" + [ testCase "effectiveInstructionAt reads signature base, override, and out-of-range" $ do + effectiveInstructionAt 0 ruled @?= ruleInstruction + effectiveInstructionAt 0 (setNodeInstr 0 "override" ruled) @?= "override" + effectiveInstructionAt 99 ruled @?= "", + testCase "setNodeInstrIfNew does not write redundant override" $ do + let kept = setNodeInstrIfNew 0 ruleInstruction ruled + changed = setNodeInstrIfNew 0 "different" ruled + instructionAt 0 kept @?= Nothing + instructionAt 0 changed @?= Just "different", + testCase "effectiveInstructionAt is index-aligned on a two-node program" $ do + effectiveInstructionAt 0 twoNode @?= "" + effectiveInstructionAt 1 twoNode @?= "Echo the sentiment label unchanged." + effectiveInstructionAt 1 (setNodeInstr 1 "changed" twoNode) @?= "changed", + testCase "tryCharge refuses an over-budget charge without changing the counter" $ do + (ok1, ok2, ok3, calls) <- + runEff . runPrim $ do + meter <- newBudgetMeter Budget {maxLmCalls = 3, maxCandidates = 32} + ok1 <- tryCharge meter 2 + ok2 <- tryCharge meter 2 + ok3 <- tryCharge meter 1 + calls <- EIORef.readIORef (meterCalls meter) + pure (ok1, ok2, ok3, calls) + (ok1, ok2, ok3, calls) @?= (True, False, True, 3), + testCase "meteredScore returns Nothing without invoking the LM when cost does not fit" $ do + ref <- newIORef (0 :: Int) + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runStubLMCounting ref $ do + meter <- newBudgetMeter Budget {maxLmCalls = 1, maxCandidates = 32} + meteredScore meter trainset exactMatch sentimentProg + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right Nothing -> pure () + Right (Just s) -> assertFailure ("expected no score, got " <> show s) + count <- readIORef ref + count @?= 0, + testCase "selectBestMetered stops at first Nothing and returns best so far" $ do + best <- + runEff . runPrim $ do + meter <- newBudgetMeter Budget {maxLmCalls = 100, maxCandidates = 32} + selectBestMetered + meter + ( \n -> + pure $ + if n >= (4 :: Int) + then Nothing + else Just (fromIntegral n) + ) + [1 :: Int, 3, 2, 4, 99] + best @?= Just (Scored 3 3.0), + testCase "withLmCallCount counts one completion per scored example" $ do + res <- runStub $ withLmCallCount (scoreOn trainset exactMatch sentimentProg) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (_score, count) -> + assertBool ("expected two completions, got " <> show count) (count == 2) + ] diff --git a/test/SeedingSpec.hs b/test/SeedingSpec.hs new file mode 100644 --- /dev/null +++ b/test/SeedingSpec.hs @@ -0,0 +1,164 @@ +{-# LANGUAGE TypeApplications #-} + +-- | Regression tests for optimizer instruction seeding. +module SeedingSpec (tests) where + +import Control.Lens ((&), (.~)) +import Data.Generics.Labels () +import Data.IORef (newIORef, readIORef) +import Data.Text qualified as T +import Effectful (Eff, IOE, runEff) +import Effectful.Concurrent (Concurrent, runConcurrent) +import Effectful.Error.Static (Error, runErrorNoCallStack) +import Effectful.Prim (Prim, runPrim) +import Shikumi.Compile.Types (compiledProgram) +import Shikumi.Effect.Time (Time, runTime) +import Shikumi.Error (ShikumiError) +import Shikumi.Eval (Dataset, dataset, exactMatch, example) +import Shikumi.Eval qualified as Eval +import Shikumi.LLM (LLM) +import Shikumi.Optimize + ( Budget (..), + CoproConfig (..), + Miprov2Auto (..), + Miprov2Config (..), + bootstrapDemoCandidates, + copro, + defaultBudget, + gepa, + instructionAt, + instructionSearch, + miprov2Auto, + miprov2With, + optimize, + recoverDemo, + reflectiveProposer, + scoreOn, + withDemos, + ) +import Shikumi.Program (Demo) +import StubLM (Label (..), Sentence (..), ruleInstruction, ruled, runGepaStubLMCapturing, runJointStubLM, runStubLM) +import Test.Tasty (TestTree, testGroup) +import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=)) + +runStub :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runStub act = runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runStubLM act + +runJoint :: Eff '[LLM, Error ShikumiError, Concurrent, Time, Prim, IOE] a -> IO (Either ShikumiError a) +runJoint act = + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ runJointStubLM act + +trainset :: Dataset Sentence Label +trainset = + dataset + [ example (Sentence "good film") (Label "positive"), + example (Sentence "good book") (Label "positive"), + example (Sentence "bad film") (Label "negative"), + example (Sentence "bad book") (Label "negative") + ] + +heldout :: Dataset Sentence Label +heldout = + dataset + [ example (Sentence "good movie") (Label "positive"), + example (Sentence "bad movie") (Label "negative") + ] + +jointTrain :: Dataset Sentence Label +jointTrain = + dataset + [ example (Sentence "good movie") (Label "positive"), + example (Sentence "bad movie") (Label "negative"), + example (Sentence "great film") (Label "positive"), + example (Sentence "terrible film") (Label "negative") + ] + +jointHeldout :: Dataset Sentence Label +jointHeldout = + dataset + [ example (Sentence "good show") (Label "positive"), + example (Sentence "bad show") (Label "negative"), + example (Sentence "great show") (Label "positive"), + example (Sentence "terrible show") (Label "negative") + ] + +coveringDemos :: [Demo] +coveringDemos = + [ recoverDemo (Sentence "great film") (Label "positive"), + recoverDemo (Sentence "terrible film") (Label "negative") + ] + +tests :: TestTree +tests = + testGroup + "instruction seeding" + [ testCase "instructionSearch never degrades a solved signature instruction" $ do + res <- + runStub $ do + before <- scoreOn heldout exactMatch ruled + cp <- optimize (instructionSearch 1 defaultBudget) trainset exactMatch ruled + after <- scoreOn heldout exactMatch (compiledProgram cp) + pure (before, after, instructionAt 0 (compiledProgram cp)) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after, override) -> do + before @?= 1.0 + assertBool ("instructionSearch: expected " <> show after <> " >= " <> show before) (after >= before) + override @?= Nothing, + testCase "copro never degrades a solved student when proposer calls are unaffordable" $ do + let cfg = CoproConfig {breadth = 2, depth = 1, budget = Budget {maxLmCalls = 4, maxCandidates = 32}} + res <- + runStub $ do + before <- scoreOn heldout exactMatch ruled + cp <- optimize (copro cfg) trainset exactMatch ruled + after <- scoreOn heldout exactMatch (compiledProgram cp) + pure (before, after, instructionAt 0 (compiledProgram cp)) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after, override) -> do + before @?= 1.0 + assertBool ("copro: expected " <> show after <> " >= " <> show before) (after >= before) + override @?= Nothing, + testCase "miprov2 never returns worse than its solved input baseline" $ do + let student = withDemos coveringDemos ruled + cfg :: Miprov2Config + cfg = miprov2Auto Miprov2Light & #budget .~ Budget {maxLmCalls = 4, maxCandidates = 32} + res <- + runJoint $ do + before <- scoreOn jointHeldout exactMatch student + cp <- optimize (miprov2With cfg student) jointTrain exactMatch student + after <- scoreOn jointHeldout exactMatch (compiledProgram cp) + pure (before, after) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right (before, after) -> do + before @?= 1.0 + assertBool ("miprov2: expected " <> show after <> " >= " <> show before) (after >= before), + testCase "miprov2 demo candidate 0 is the node's current demos" $ do + let student = withDemos coveringDemos ruled + cfg :: Miprov2Config + cfg = miprov2Auto Miprov2Light & #numDemoCandidates .~ 3 + res <- runJoint (bootstrapDemoCandidates cfg student jointTrain exactMatch student) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right [sets] -> do + take 1 sets @?= [coveringDemos] + assertBool "empty demo set remains available" ([] `elem` sets) + Right other -> assertFailure ("expected one node, got " <> show (length other)), + testCase "gepa reflection prompt sees the effective signature instruction" $ do + ref <- newIORef [] + let fbMetric _ _ = (Eval.boolScore False, "be more specific") + res <- + runEff . runPrim . runTime . runConcurrent . runErrorNoCallStack @ShikumiError $ + runGepaStubLMCapturing ref (optimize (gepa reflectiveProposer fbMetric defaultBudget) trainset exactMatch ruled) + case res of + Left e -> assertFailure ("unexpected error: " <> show e) + Right _ -> pure () + captured <- readIORef ref + case filter ("proposedInstruction" `T.isInfixOf`) captured of + firstReflection : _ -> + assertBool + "expected first GEPA reflection request to include the effective RULE instruction" + (("currentInstruction: " <> ruleInstruction) `T.isInfixOf` firstReflection) + [] -> assertFailure "expected at least one GEPA reflection request" + ] diff --git a/test/StubLM.hs b/test/StubLM.hs --- a/test/StubLM.hs +++ b/test/StubLM.hs @@ -39,6 +39,9 @@ -- * Signature and program sentimentSig, sentimentProg, + ruled, + echoSig, + sentimentPipeline, -- * Ground truth and helpers goldLabel, @@ -56,6 +59,7 @@ -- * The reflective task (EP-22, GEPA) runGepaStubLM, + runGepaStubLMCapturing, ) where @@ -82,6 +86,7 @@ import Effectful.Dispatch.Dynamic (interpret) import GHC.Generics (Generic) import Shikumi.Adapter (ToPrompt) +import Shikumi.Combinator ((>>>)) import Shikumi.LLM (LLM (..)) import Shikumi.Module (predict) import Shikumi.Program (Program) @@ -132,6 +137,19 @@ sentimentProg :: Program Sentence Label sentimentProg = predict sentimentSig +-- | A single-node sentiment program whose signature instruction already solves the +-- task under the stub. Optimizers must preserve it when "keep current" wins. +ruled :: Program Sentence Label +ruled = predict (mkSignature ruleInstruction) + +-- | A second-stage identity node for multi-node optimizer tests. +echoSig :: Signature Label Label +echoSig = mkSignature "Echo the sentiment label unchanged." + +-- | A two-node sentiment pipeline: classify a sentence, then echo the label. +sentimentPipeline :: Program Sentence Label +sentimentPipeline = sentimentProg >>> predict echoSig + -- --------------------------------------------------------------------------- -- Ground truth and helpers -- --------------------------------------------------------------------------- @@ -206,6 +224,7 @@ -- | The joint task's classification rule (see 'runJointStubLM'). answerJoint :: Context -> Text answerJoint ctx + | Just lbl <- parseEcho (lastUserText ctx) = lbl | "good" `elem` ws || "bad" `elem` ws = -- region A: correct only with a RULE instruction if instructionHasRule ctx then goldLabel s else "neutral" @@ -262,6 +281,15 @@ Complete _ ctx _ -> pure (mkResponse (respondGepa ctx)) Stream {} -> pure [] +-- | Like 'runGepaStubLM' but records each request's rendered text for prompt-signal +-- assertions. +runGepaStubLMCapturing :: (IOE :> es) => IORef [Text] -> Eff (LLM : es) a -> Eff es a +runGepaStubLMCapturing ref = interpret $ \_ -> \case + Complete _ ctx _ -> do + liftIO (modifyIORef' ref (++ [fullRequestText ctx])) + pure (mkResponse (respondGepa ctx)) + Stream {} -> pure [] + respondGepa :: Context -> Text respondGepa ctx | sysHasMarker "proposedInstruction" ctx = @@ -283,9 +311,11 @@ -- | Classify the actual input given the demos and instruction in the context. answerSentiment :: Context -> Text answerSentiment ctx = - case demos of - [] -> if instructionHasRule ctx then goldLabel s else "neutral" - _ -> nnLabel s demos + case parseEcho (lastUserText ctx) of + Just lbl -> lbl + Nothing -> case demos of + [] -> if instructionHasRule ctx then goldLabel s else "neutral" + _ -> nnLabel s demos where msgs = V.toList (ctx ^. #messages) demos = demoPairs msgs @@ -364,6 +394,10 @@ parseSentence t = T.strip (fromMaybe stripped (T.stripPrefix "text:" stripped)) where stripped = T.strip t + +-- | Parse the rendered input to the echo node, whose input type is 'Label'. +parseEcho :: Text -> Maybe Text +parseEcho t = T.strip <$> T.stripPrefix "sentiment:" (T.strip t) -- | Read the @sentiment@ marker section out of a rendered demo output. parseLabel :: Text -> Text