diff --git a/src/benchmarks/README.md b/src/benchmarks/README.md index 8d43cf9..46d634e 100644 --- a/src/benchmarks/README.md +++ b/src/benchmarks/README.md @@ -36,6 +36,7 @@ interface Benchmark { | `locomo` | GitHub snap-research/locomo | Long context memory benchmark | | `longmemeval` | HuggingFace xiaowu0162/longmemeval-cleaned | Long-term memory evaluation | | `convomem` | HuggingFace Salesforce/ConvoMem | Conversational memory benchmark | +| `infinitebench` | HuggingFace xinrongzhang2022/InfiniteBench | Long-context understanding benchmark for testing retrieval, reasoning, and comprehension over extremely long inputs | ## Question Types diff --git a/src/benchmarks/index.ts b/src/benchmarks/index.ts index b790e87..f05e06c 100644 --- a/src/benchmarks/index.ts +++ b/src/benchmarks/index.ts @@ -2,11 +2,13 @@ import type { Benchmark, BenchmarkName } from "../types/benchmark" import { LoCoMoBenchmark } from "./locomo" import { LongMemEvalBenchmark } from "./longmemeval" import { ConvoMemBenchmark } from "./convomem" +import { InfiniteBenchBenchmark } from "./infinitebench" const benchmarks: Record Benchmark> = { locomo: LoCoMoBenchmark, longmemeval: LongMemEvalBenchmark, convomem: ConvoMemBenchmark, + infinitebench: InfiniteBenchBenchmark, } export function createBenchmark(name: BenchmarkName): Benchmark { @@ -21,4 +23,4 @@ export function getAvailableBenchmarks(): BenchmarkName[] { return Object.keys(benchmarks) as BenchmarkName[] } -export { LoCoMoBenchmark, LongMemEvalBenchmark, ConvoMemBenchmark } +export { LoCoMoBenchmark, LongMemEvalBenchmark, ConvoMemBenchmark, InfiniteBenchBenchmark } diff --git a/src/benchmarks/infinitebench/index.ts b/src/benchmarks/infinitebench/index.ts new file mode 100644 index 0000000..56a5680 --- /dev/null +++ b/src/benchmarks/infinitebench/index.ts @@ -0,0 +1,162 @@ +import { existsSync, readFileSync } from "fs" +import { join } from "path" + +import type { Benchmark, BenchmarkConfig, QuestionFilter } from "../../types/benchmark" +import type { + UnifiedMessage, + UnifiedQuestion, + UnifiedSession, + QuestionTypeRegistry, +} from "../../types/unified" + +import type { InfiniteBenchItem } from "./types" +import { logger } from "../../utils/logger" + +const DEFAULT_DATA_PATH = "./data/benchmarks/infinitebench" + +const TASKS = [ + "passkey", + "kv_retrieval", + "number_string", + "code_run", + "code_debug", + "math_find", + "math_calc", + "longdialogue_qa_eng", + "longbook_qa_eng", + "longbook_sum_eng", + "longbook_choice_eng", + "longbook_qa_chn", +] as const + +type InfiniteBenchTask = (typeof TASKS)[number] + +export const INFINITEBENCH_QUESTION_TYPES: QuestionTypeRegistry = Object.fromEntries( + TASKS.map((task) => [ + task, + { + id: task, + alias: task, + description: `InfiniteBench ${task} task`, + }, + ]) +) + +export class InfiniteBenchBenchmark implements Benchmark { + name = "infinitebench" + + private questions: UnifiedQuestion[] = [] + private sessions: Map = new Map() + + private dataPath = DEFAULT_DATA_PATH + + async load(config?: BenchmarkConfig): Promise { + this.dataPath = config?.dataPath ?? DEFAULT_DATA_PATH + + const datasetPath = join(process.cwd(), this.dataPath) + + if (!existsSync(datasetPath)) { + throw new Error(`InfiniteBench dataset missing: ${datasetPath}`) + } + + this.loadDataset(datasetPath) + } + + private loadDataset(datasetPath: string): void { + for (const task of TASKS) { + const file = join(datasetPath, `${task}.jsonl`) + + if (!existsSync(file)) { + logger.warn(`Skipping missing task file: ${file}`) + continue + } + + const rows = readFileSync(file, "utf-8").split("\n").filter(Boolean) + + for (const row of rows) { + try { + const item: InfiniteBenchItem = JSON.parse(row) + + this.addQuestion(item, task) + } catch (error) { + logger.error(`Failed parsing ${file}: ${error}`) + } + } + } + + logger.info(`Loaded ${this.questions.length} InfiniteBench questions`) + } + + private addQuestion(item: InfiniteBenchItem, task: InfiniteBenchTask): void { + const questionId = `infinitebench-${task}-${item.id}` + + const session = this.createSession(item.context, questionId) + + const question: UnifiedQuestion = { + questionId, + + question: item.input, + + questionType: task, + + groundTruth: item.answer.at(0) ?? "", + + haystackSessionIds: [session.sessionId], + + metadata: { + task, + options: item.options, + }, + } + + this.questions.push(question) + + this.sessions.set(questionId, [session]) + } + + private createSession(context: string, questionId: string): UnifiedSession { + const message: UnifiedMessage = { + role: "user", + + content: context, + } + + return { + sessionId: `${questionId}-session`, + + messages: [message], + } + } + + getQuestions(filter?: QuestionFilter): UnifiedQuestion[] { + let result = this.questions + + if (filter?.questionTypes?.length) { + result = result.filter((q) => filter.questionTypes!.includes(q.questionType)) + } + + if (filter?.offset) { + result = result.slice(filter.offset) + } + + if (filter?.limit) { + result = result.slice(0, filter.limit) + } + + return result + } + + getHaystackSessions(questionId: string): UnifiedSession[] { + return this.sessions.get(questionId) ?? [] + } + + getGroundTruth(questionId: string): string { + return this.questions.find((q) => q.questionId === questionId)?.groundTruth ?? "" + } + + getQuestionTypes(): QuestionTypeRegistry { + return INFINITEBENCH_QUESTION_TYPES + } +} + +export default InfiniteBenchBenchmark \ No newline at end of file diff --git a/src/benchmarks/infinitebench/types.ts b/src/benchmarks/infinitebench/types.ts new file mode 100644 index 0000000..12139f8 --- /dev/null +++ b/src/benchmarks/infinitebench/types.ts @@ -0,0 +1,7 @@ +export interface InfiniteBenchItem { + id: number + context: string + input: string + answer: string[] + options: string[] +} \ No newline at end of file diff --git a/src/cli/index.ts b/src/cli/index.ts index b3c29d6..4fb20eb 100644 --- a/src/cli/index.ts +++ b/src/cli/index.ts @@ -150,11 +150,16 @@ Available benchmark datasets for evaluation: convomem ConvoMem - Conversational memory benchmark Tests: user facts, assistant facts, preferences, implicit connections Source: HuggingFace Salesforce/ConvoMem (downloaded on first use) + + infinitebench InfiniteBench - Ultra-long context understanding benchmark + Tests: document QA, long-context retrieval, code understanding, multi-hop reasoning + Source: GitHub OpenBMB/InfiniteBench (downloaded on first use) Usage: -b locomo Run LoCoMo benchmark -b longmemeval Run LongMemEval benchmark -b convomem Run ConvoMem benchmark + -b infinitebench Run InfiniteBench benchmark `) } diff --git a/src/types/benchmark.ts b/src/types/benchmark.ts index 07961e4..c73b384 100644 --- a/src/types/benchmark.ts +++ b/src/types/benchmark.ts @@ -20,4 +20,4 @@ export interface Benchmark { getQuestionTypes(): QuestionTypeRegistry } -export type BenchmarkName = "locomo" | "longmemeval" | "convomem" +export type BenchmarkName = "locomo" | "longmemeval" | "convomem" | "infinitebench"