Skip to content

fix(sampling): use an unbiased, seedable shuffle for random question selection #86

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Agnik47 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(sampling): use an unbiased, seedable shuffle for random question selection #86

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 12 additions & 0 deletions src/cli/commands/compare.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ interface CompareArgs {
compareId?: string
sample?: number
sampleType?: SampleType
seed?: number
limit?: number
force?: boolean
}
Expand Down Expand Up @@ -50,6 +51,13 @@ export function parseCompareArgs(args: string[]): CompareArgs | null {
logger.error(`Invalid sample type: ${type}. Valid types: consecutive, random`)
return null
}
} else if (arg === "--seed") {
const seed = Number(args[++i])
if (!Number.isInteger(seed) || seed < 0) {
logger.error(`Invalid seed: ${seed}. Expected a non-negative integer.`)
return null
}
parsed.seed = seed
} else if (arg === "-l" || arg === "--limit") {
parsed.limit = parseInt(args[++i], 10)
} else if (arg === "--force") {
Expand Down Expand Up @@ -83,6 +91,9 @@ export async function compareCommand(args: string[]): Promise<void> {
console.log(` -m, --answering-model Answering model (default: ${DEFAULT_ANSWERING_MODEL})`)
console.log(" -s, --sample Sample N questions per category")
console.log(" --sample-type Sample type: consecutive (default), random")
console.log(
" --seed N Seed for --sample-type random (logged and stored when omitted)"
)
console.log(" -l, --limit Limit total number of questions")
console.log(" --compare-id Compare ID (for resuming)")
console.log(" --force Clear existing comparison and start fresh")
Expand Down Expand Up @@ -122,6 +133,7 @@ export async function compareCommand(args: string[]): Promise<void> {
mode: "sample",
sampleType: parsed.sampleType || "consecutive",
perCategory: parsed.sample,
seed: parsed.seed,
}
} else if (parsed.limit) {
sampling = {
Expand Down
12 changes: 12 additions & 0 deletions src/cli/commands/run.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ interface RunArgs {
limit?: number
sample?: number
sampleType?: SampleType
seed?: number
force?: boolean
fromPhase?: PhaseId
concurrency?: ConcurrencyConfig
Expand Down Expand Up @@ -60,6 +61,13 @@ export function parseRunArgs(args: string[]): RunArgs | null {
logger.error(`Invalid sample type: ${type}. Valid types: consecutive, random`)
return null
}
} else if (arg === "--seed") {
const seed = Number(args[++i])
if (!Number.isInteger(seed) || seed < 0) {
logger.error(`Invalid seed: ${seed}. Expected a non-negative integer.`)
return null
}
parsed.seed = seed
} else if (arg === "-f" || arg === "--from-phase") {
const phase = args[++i] as PhaseId
if (PHASE_ORDER.includes(phase)) {
Expand Down Expand Up @@ -121,6 +129,9 @@ export async function runCommand(args: string[]): Promise<void> {
console.log(` -m, --answering-model Answering model (default: ${DEFAULT_ANSWERING_MODEL})`)
console.log(" -s, --sample Sample N questions per category")
console.log(" --sample-type Sample type: consecutive (default), random")
console.log(
" --seed N Seed for --sample-type random (logged and stored when omitted)"
)
console.log(" -l, --limit Limit total number of questions to process")
console.log(` -f, --from-phase Start from phase: ${PHASE_ORDER.join(", ")}`)
console.log(" --concurrency N Default concurrency for all phases")
Expand Down Expand Up @@ -191,6 +202,7 @@ export async function runCommand(args: string[]): Promise<void> {
mode: "sample",
sampleType: parsed.sampleType || "consecutive",
perCategory: parsed.sample,
seed: parsed.seed,
}
} else if (parsed.limit) {
sampling = {
Expand Down
44 changes: 12 additions & 32 deletions src/orchestrator/batch.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import type { BenchmarkName } from "../types/benchmark"
import type { SamplingConfig } from "../types/checkpoint"
import type { BenchmarkResult } from "../types/unified"
import { orchestrator, CheckpointManager } from "./index"
import { selectQuestionsBySampling } from "./sampling"
import { createBenchmark } from "../benchmarks"
import { logger } from "../utils/logger"
import { existsSync, mkdirSync, readFileSync, writeFileSync, rmSync } from "fs"
Expand Down Expand Up @@ -52,36 +53,6 @@ function generateCompareId(): string {
return `compare-${date}-${time}`
}

function selectQuestionsBySampling(
allQuestions: { questionId: string; questionType: string }[],
sampling: SamplingConfig
): string[] {
if (sampling.mode === "full") {
return allQuestions.map((q) => q.questionId)
}
if (sampling.mode === "limit" && sampling.limit) {
return allQuestions.slice(0, sampling.limit).map((q) => q.questionId)
}
if (sampling.mode === "sample" && sampling.perCategory) {
const byType: Record<string, { questionId: string; questionType: string }[]> = {}
for (const q of allQuestions) {
if (!byType[q.questionType]) byType[q.questionType] = []
byType[q.questionType].push(q)
}
const selected: string[] = []
for (const questions of Object.values(byType)) {
if (sampling.sampleType === "random") {
const shuffled = [...questions].sort(() => Math.random() - 0.5)
selected.push(...shuffled.slice(0, sampling.perCategory).map((q) => q.questionId))
} else {
selected.push(...questions.slice(0, sampling.perCategory).map((q) => q.questionId))
}
}
return selected
}
return allQuestions.map((q) => q.questionId)
}

export class BatchManager {
private getComparePath(compareId: string): string {
return join(COMPARE_DIR, compareId)
Expand Down Expand Up @@ -155,8 +126,17 @@ export class BatchManager {
const allQuestions = benchmarkInstance.getQuestions()

let targetQuestionIds: string[]
let resolvedSampling = sampling
if (sampling) {
targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling)
const selection = selectQuestionsBySampling(allQuestions, sampling)
targetQuestionIds = selection.questionIds
// Persist the seed that was actually used so the comparison can be rebuilt
// on the same subset later.
resolvedSampling =
selection.seed === undefined ? sampling : { ...sampling, seed: selection.seed }
if (selection.seed !== undefined) {
logger.info(`Random sampling seed: ${selection.seed} (reuse with --seed ${selection.seed})`)
}
} else {
targetQuestionIds = allQuestions.map((q) => q.questionId)
}
Expand All @@ -168,7 +148,7 @@ export class BatchManager {
benchmark,
judge: judgeModel,
answeringModel,
sampling,
sampling: resolvedSampling,
targetQuestionIds,
runs: providers.map((provider) => ({
provider,
Expand Down
48 changes: 12 additions & 36 deletions src/orchestrator/index.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import { createProvider } from "../providers"
import { createBenchmark } from "../benchmarks"
import { createJudge } from "../judges"
import { CheckpointManager } from "./checkpoint"
import { selectQuestionsBySampling } from "./sampling"
import { getProviderConfig, getJudgeConfig } from "../utils/config"
import { resolveModel } from "../utils/models"
import { logger } from "../utils/logger"
Expand All @@ -31,40 +32,6 @@ export interface OrchestratorOptions {
phases?: ("ingest" | "indexing" | "search" | "answer" | "evaluate" | "report")[]
}

function selectQuestionsBySampling(
allQuestions: { questionId: string; questionType: string }[],
sampling: SamplingConfig
): string[] {
if (sampling.mode === "full") {
return allQuestions.map((q) => q.questionId)
}

if (sampling.mode === "limit" && sampling.limit) {
return allQuestions.slice(0, sampling.limit).map((q) => q.questionId)
}

if (sampling.mode === "sample" && sampling.perCategory) {
const byType: Record<string, { questionId: string; questionType: string }[]> = {}
for (const q of allQuestions) {
if (!byType[q.questionType]) byType[q.questionType] = []
byType[q.questionType].push(q)
}

const selected: string[] = []
for (const questions of Object.values(byType)) {
if (sampling.sampleType === "random") {
const shuffled = [...questions].sort(() => Math.random() - 0.5)
selected.push(...shuffled.slice(0, sampling.perCategory).map((q) => q.questionId))
} else {
selected.push(...questions.slice(0, sampling.perCategory).map((q) => q.questionId))
}
}
return selected
}

return allQuestions.map((q) => q.questionId)
}

export class Orchestrator {
private checkpointManager: CheckpointManager

Expand Down Expand Up @@ -217,11 +184,20 @@ export class Orchestrator {
targetQuestionIds = questionIds
} else if (sampling) {
logger.info(`Using sampling mode: ${sampling.mode}`)
targetQuestionIds = selectQuestionsBySampling(allQuestions, sampling)
checkpoint.sampling = sampling
const selection = selectQuestionsBySampling(allQuestions, sampling)
targetQuestionIds = selection.questionIds
// Record the resolved seed, not the requested config, so the checkpoint
// is enough to reproduce this exact subset.
checkpoint.sampling =
selection.seed === undefined ? sampling : { ...sampling, seed: selection.seed }
logger.info(
`Sampling selected ${targetQuestionIds.length} questions from ${allQuestions.length} total`
)
if (selection.seed !== undefined) {
logger.info(
`Random sampling seed: ${selection.seed} (reuse with --seed ${selection.seed})`
)
}
} else if (effectiveLimit) {
logger.info(`Using limit: ${effectiveLimit}`)
targetQuestionIds = allQuestions.slice(0, effectiveLimit).map((q) => q.questionId)
Expand Down
Loading