[Feature] Upgrade the needle-in-a-haystack experiment to Needlebench (#913)

* add needlebench

* simplify needlebench 32k, 128k, 200k for eval

* update act prompt

* fix bug in needlebench summarizer

* add needlebench intro, fix summarizer

* lint summarizer

* fix linting error

* move readme.md

* update readme for needlebench

* update docs of needlebench

* simplify needlebench summarizers

configs/datasets/needlebench/atc/atc.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.atc import NeedleBenchATCOrderedDataset
+from opencompass.datasets.needlebench.atc import NeedleBenchATCDataset
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchOriginEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+needle_num_list = list(range(2, 100, 3))
+document_depth_percent_intervals = 20
+repeats = 30
+names_path = './data/needlebench/names.json'
+
+needlebench_atc_datasets_zh = []
+needlebench_atc_datasets_en = []
+needlebench_atc_datasets_zh_ordered = []
+needlebench_atc_datasets_en_ordered = []
+
+for num_needles in needle_num_list:
+    # ordered English version
+    dataset_dict = {
+        'abbr': f'needlebench_atc_challenge'
+                f'needle_{num_needles}_en_ordered',
+        'type': NeedleBenchATCOrderedDataset,
+        'path': names_path,
+        'num_needles': num_needles,
+        'language': 'English',
+        'repeats': repeats,
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_atc_datasets_en_ordered.append(dataset_dict)
+
+
+for num_needles in needle_num_list:
+    # ordered Chinese version
+    dataset_dict = {
+        'abbr': f'needlebench_atc_challenge'
+                f'needle_{num_needles}_zh_ordered',
+        'type': NeedleBenchATCOrderedDataset,
+        'path': names_path,
+        'num_needles': num_needles,
+        'language': 'Chinese',
+        'repeats': repeats,
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_atc_datasets_zh_ordered.append(dataset_dict)
+
+for num_needles in needle_num_list:
+    # standard English version
+    dataset_dict = {
+        'abbr': f'needlebench_atc_challenge'
+                f'needle_{num_needles}_en',
+        'type': NeedleBenchATCDataset,
+        'path': names_path,
+        'num_needles': num_needles,
+        'language': 'English',
+        'repeats': repeats,
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_atc_datasets_en.append(dataset_dict)
+
+for num_needles in needle_num_list:
+    # standard Chinese version
+    dataset_dict = {
+        'abbr': f'needlebench_atc_challenge'
+                f'needle_{num_needles}_zh',
+        'type': NeedleBenchATCDataset,
+        'path': names_path,
+        'num_needles': num_needles,
+        'language': 'Chinese',
+        'repeats': repeats,
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_atc_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_4k.needlebench import needlebench_datasets as needlebench_datasets_4k
+    from .needlebench_8k.needlebench import needlebench_datasets as needlebench_datasets_8k
+    from .needlebench_32k.needlebench import needlebench_datasets as needlebench_datasets_32k
+    from .needlebench_128k.needlebench import needlebench_datasets as needlebench_datasets_128k
+    from .needlebench_200k.needlebench import needlebench_datasets as needlebench_datasets_200k
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_128k/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_en as needlebench_multi_2needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_en as needlebench_multi_3needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_en as needlebench_multi_4needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_en as needlebench_multi_5needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_zh as needlebench_multi_2needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_zh as needlebench_multi_3needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_zh as needlebench_multi_4needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_zh as needlebench_multi_5needle_zh_datasets
+
+    from .needlebench_single import needlebench_datasets_en as needlebench_origin_en_datasets
+    from .needlebench_single import needlebench_datasets_zh as needlebench_origin_zh_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_en as needlebench_parallel_en_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_zh as needlebench_parallel_zh_datasets
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_128k/needlebench_multi_reasoning.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiDataset
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchMultiEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000])
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+# ----------English Version----------
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_en.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_en = []
+language = 'English'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_en.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_en.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_en.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_en.append(dataset_dict)
+
+# ----------Chinese Version----------
+base_path = './data/needlebench'
+file_list = ['zh_finance.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_zh.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_zh = []
+language = 'Chinese'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_zh.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_zh.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_zh.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_128k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_128k/needlebench_multi_retrieval.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelDataset
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchParallelEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000])
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+depths = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_en_128k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 3000,
+        'guide': True,
+        'language': 'English',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_en.append(dataset_dict)
+
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_zh_128k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 200,
+        'guide': True,
+        'language': 'Chinese',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_zh.append(dataset_dict)

configs/datasets/cdme/cdme8k.py → configs/datasets/needlebench/needlebench_128k/needlebench_single.py

 from opencompass.openicl.icl_prompt_template import PromptTemplate
 from opencompass.openicl.icl_retriever import ZeroRetriever
 from opencompass.openicl.icl_inferencer import GenInferencer
-from opencompass.datasets.cdme.cdme import CDMEDataset
-from opencompass.datasets.cdme.cdme import CDMEEvaluator
-from opencompass.datasets.cdme.cdme import cdme_postprocess
-from opencompass.datasets.cdme.cdme import cdme_dataset_postprocess
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginDataset
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
 import math
 
 
@@ -31,37 +31,69 @@ def generate_depth_percents(intervals, interval_type):
         raise ValueError('Unsupported interval type')
 
 
-cdme_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
 
-cdme_infer_cfg = dict(
+needlebench_infer_cfg = dict(
     prompt_template=dict(
         type=PromptTemplate,
-        template='''{prompt}'''),
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
     retriever=dict(type=ZeroRetriever),
-    inferencer=dict(type=GenInferencer, max_out_len=512))
+    inferencer=dict(type=GenInferencer))
 
-cdme_eval_cfg = dict(
-    evaluator=dict(type=CDMEEvaluator),
-    pred_postprocessor=dict(type=cdme_postprocess),
-    dataset_postprocessor=dict(type=cdme_dataset_postprocess),
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchOriginEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
     pred_role='BOT')
 
-context_lengths = list(range(1000, 9000, 1000))
+context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000])
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
 document_depth_percent_intervals = 20
 document_depth_percent_interval_type = "linear"
 
-base_path = './data/CDME'
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_en_128k',
+            'type': NeedleBenchOriginDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': 'English',
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_en.append(dataset_dict)
+
 file_list = ['zh_finance.jsonl']
-cdme_datasets = []
+needlebench_datasets_zh = []
+needle_file_name = 'needles.jsonl'
 
 for original_context_length in context_lengths:
-    for depth_percent in generate_depth_percents(
-            document_depth_percent_intervals,
-            document_depth_percent_interval_type):
+    for depth_percent in depths_list:
         dataset_dict = {
-            'abbr': f'CDME_Length{original_context_length}'
-            f'Depth{int(depth_percent)}',
-            'type': CDMEDataset,
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_zh_128k',
+            'type': NeedleBenchOriginDataset,
             'path': base_path,
             'length': original_context_length,
             'depth': int(depth_percent),
@@ -71,11 +103,9 @@ for original_context_length in context_lengths:
             'length_buffer': 200,
             'guide': True,
             'language': 'Chinese',
-            'needle': '\n小明最喜欢的实习的地点就是上海人工智能实验室。\n',
-            'retrieval_question': '小明最喜欢的实习地点是哪里？请按照'
-            '“小明最喜欢的实习地点就是________。”的格式回答。',
-            'reader_cfg': cdme_reader_cfg,
-            'infer_cfg': cdme_infer_cfg,
-            'eval_cfg': cdme_eval_cfg
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
         }
-        cdme_datasets.append(dataset_dict)
+        needlebench_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_200k/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_en as needlebench_multi_2needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_en as needlebench_multi_3needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_en as needlebench_multi_4needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_en as needlebench_multi_5needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_zh as needlebench_multi_2needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_zh as needlebench_multi_3needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_zh as needlebench_multi_4needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_zh as needlebench_multi_5needle_zh_datasets
+
+    from .needlebench_single import needlebench_datasets_en as needlebench_origin_en_datasets
+    from .needlebench_single import needlebench_datasets_zh as needlebench_origin_zh_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_en as needlebench_parallel_en_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_zh as needlebench_parallel_zh_datasets
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_200k/needlebench_multi_reasoning.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiDataset
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchMultiEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+# context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000, 144000, 160000, 176000, 192000, 200000])
+context_lengths = [16000, 48000, 80000, 112000, 128000, 144000, 176000, 200000]
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+
+# ----------English Version----------
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_en.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_en = []
+language = 'English'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_en.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_en.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_en.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_en.append(dataset_dict)
+
+# ----------Chinese Version----------
+base_path = './data/needlebench'
+file_list = ['zh_finance.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_zh.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_zh = []
+language = 'Chinese'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_zh.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_zh.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_zh.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_200k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_200k/needlebench_multi_retrieval.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelDataset
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchParallelEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+# context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000, 144000, 160000, 176000, 192000, 200000])
+context_lengths = list([16000, 48000, 80000, 112000, 128000, 144000, 176000, 200000])
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+depths = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_en_200k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 3000,
+        'guide': True,
+        'language': 'English',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_en.append(dataset_dict)
+
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_zh_200k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 200,
+        'guide': True,
+        'language': 'Chinese',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_200k/needlebench_single.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginDataset
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchOriginEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+# context_lengths = list([16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000, 144000, 160000, 176000, 192000, 200000])
+context_lengths = [16000, 48000, 80000, 112000, 128000, 144000, 176000, 200000]
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_en_200k',
+            'type': NeedleBenchOriginDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': 'English',
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_en.append(dataset_dict)
+
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+needle_file_name = 'needles.jsonl'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_zh_200k',
+            'type': NeedleBenchOriginDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': 'Chinese',
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_32k/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_en as needlebench_multi_2needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_en as needlebench_multi_3needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_en as needlebench_multi_4needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_en as needlebench_multi_5needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_zh as needlebench_multi_2needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_zh as needlebench_multi_3needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_zh as needlebench_multi_4needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_zh as needlebench_multi_5needle_zh_datasets
+
+    from .needlebench_single import needlebench_datasets_en as needlebench_origin_en_datasets
+    from .needlebench_single import needlebench_datasets_zh as needlebench_origin_zh_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_en as needlebench_parallel_en_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_zh as needlebench_parallel_zh_datasets
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_32k/needlebench_multi_reasoning.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiDataset
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchMultiEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list([9000, 13000, 17000, 21000, 25000, 29000, 31000, 32000])
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+# ----------English Version----------
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_en.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_en = []
+language = 'English'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 3000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_en.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 3000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_en.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 3000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_en.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 3000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_en.append(dataset_dict)
+
+# ----------Chinese Version----------
+base_path = './data/needlebench'
+file_list = ['zh_finance.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_zh.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_zh = []
+language = 'Chinese'
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_zh.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_zh.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_zh.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_32k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_32k/needlebench_multi_retrieval.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelDataset
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchParallelEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list([9000, 13000, 17000, 21000, 25000, 29000, 31000, 32000])
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+depths = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_en_32k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 3000,
+        'guide': True,
+        'language': 'English',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_en.append(dataset_dict)
+
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_zh_32k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 200,
+        'guide': True,
+        'language': 'Chinese',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_zh.append(dataset_dict)

configs/datasets/cdme/multi_needle/cdme8k_cot2_italy.py → configs/datasets/needlebench/needlebench_32k/needlebench_single.py

 from opencompass.openicl.icl_prompt_template import PromptTemplate
 from opencompass.openicl.icl_retriever import ZeroRetriever
 from opencompass.openicl.icl_inferencer import GenInferencer
-from opencompass.datasets.cdme.cdme_multi import CDMEDataset
-from opencompass.datasets.cdme.cdme_multi import CDMEEvaluator
-from opencompass.datasets.cdme.cdme_multi import cdme_postprocess
-from opencompass.datasets.cdme.cdme_multi import cdme_dataset_postprocess
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginDataset
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
 import math
 
 
@@ -31,46 +31,69 @@ def generate_depth_percents(intervals, interval_type):
         raise ValueError('Unsupported interval type')
 
 
-cdme_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
 
-cdme_infer_cfg = dict(
+needlebench_infer_cfg = dict(
     prompt_template=dict(
         type=PromptTemplate,
-        template='''{prompt}'''),
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
     retriever=dict(type=ZeroRetriever),
-    inferencer=dict(type=GenInferencer, max_out_len=512))
+    inferencer=dict(type=GenInferencer))
 
-cdme_eval_cfg = dict(
-    evaluator=dict(type=CDMEEvaluator),
-    pred_postprocessor=dict(type=cdme_postprocess),
-    dataset_postprocessor=dict(type=cdme_dataset_postprocess),
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchOriginEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
     pred_role='BOT')
 
-context_lengths = list(range(1000, 9000, 1000))
+context_lengths = list([9000, 13000, 17000, 21000, 25000, 29000, 31000, 32000])
+depths_list = [0, 10, 21, 31, 42, 52, 63, 73, 84, 94, 100]
 document_depth_percent_intervals = 20
 document_depth_percent_interval_type = "linear"
 
-base_path = './data/CDME'
-file_list = ['zh_finance.jsonl']
-cdme_datasets = []
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
 
-needles = ['\n意大利的佛罗伦萨有一家名为"La Giostra"的餐馆，是整个佛罗伦萨中排行第一的餐馆。\n',
-           '"La Giostra"餐馆的特色菜肴是松露奶酪通心粉。']
-retrieval_question = ("佛罗伦萨中排行第一的餐馆的特色菜肴是？"
-                      "请按照'佛罗伦萨中排行第一的餐馆的特色菜肴是______。'的格式回答。")
-answer = "佛罗伦萨中排行第一的餐馆的特色菜肴是松露奶酪通心粉。"
+for original_context_length in context_lengths:
+    for depth_percent in depths_list:
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_en_32k',
+            'type': NeedleBenchOriginDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 3000,
+            'guide': True,
+            'language': 'English',
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_en.append(dataset_dict)
 
-keyword = "松露奶酪通心粉"
-diff = 25
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+needle_file_name = 'needles.jsonl'
 
 for original_context_length in context_lengths:
-    for depth_percent in generate_depth_percents(
-            document_depth_percent_intervals,
-            document_depth_percent_interval_type):
+    for depth_percent in depths_list:
         dataset_dict = {
-            'abbr': f'CDME_Length{original_context_length}'
-            f'Depth{int(depth_percent)}',
-            'type': CDMEDataset,
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_zh_32k',
+            'type': NeedleBenchOriginDataset,
             'path': base_path,
             'length': original_context_length,
             'depth': int(depth_percent),
@@ -80,13 +103,9 @@ for original_context_length in context_lengths:
             'length_buffer': 200,
             'guide': True,
             'language': 'Chinese',
-            'needles': needles,
-            'diff': diff,
-            'retrieval_question': retrieval_question,
-            'answer': answer,
-            'keyword': keyword,
-            'reader_cfg': cdme_reader_cfg,
-            'infer_cfg': cdme_infer_cfg,
-            'eval_cfg': cdme_eval_cfg
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
         }
-        cdme_datasets.append(dataset_dict)
+        needlebench_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_4k/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_en as needlebench_multi_2needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_en as needlebench_multi_3needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_en as needlebench_multi_4needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_en as needlebench_multi_5needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_zh as needlebench_multi_2needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_zh as needlebench_multi_3needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_zh as needlebench_multi_4needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_zh as needlebench_multi_5needle_zh_datasets
+
+    from .needlebench_single import needlebench_datasets_en as needlebench_origin_en_datasets
+    from .needlebench_single import needlebench_datasets_zh as needlebench_origin_zh_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_en as needlebench_parallel_en_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_zh as needlebench_parallel_zh_datasets
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_4k/needlebench_multi_reasoning.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiDataset
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchMultiEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list(range(1000, 5000, 1000))
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+# ----------English Version----------
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_en.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_en = []
+language = 'English'
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_en.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_en.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_en.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 600,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_en.append(dataset_dict)
+
+# ----------Chinese Version----------
+base_path = './data/needlebench'
+file_list = ['zh_finance.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_zh.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_zh = []
+language = 'Chinese'
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_zh.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_zh.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_zh.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_4k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_zh.append(dataset_dict)

configs/datasets/cdme/cdme32k.py → configs/datasets/needlebench/needlebench_4k/needlebench_multi_retrieval.py

 from opencompass.openicl.icl_prompt_template import PromptTemplate
 from opencompass.openicl.icl_retriever import ZeroRetriever
 from opencompass.openicl.icl_inferencer import GenInferencer
-from opencompass.datasets.cdme.cdme import CDMEDataset
-from opencompass.datasets.cdme.cdme import CDMEEvaluator
-from opencompass.datasets.cdme.cdme import cdme_postprocess
-from opencompass.datasets.cdme.cdme import cdme_dataset_postprocess
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelDataset
+from opencompass.datasets.needlebench.parallel import NeedleBenchParallelEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
 import math
 
 
@@ -31,51 +31,81 @@ def generate_depth_percents(intervals, interval_type):
         raise ValueError('Unsupported interval type')
 
 
-cdme_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
 
-cdme_infer_cfg = dict(
+needlebench_infer_cfg = dict(
     prompt_template=dict(
         type=PromptTemplate,
-        template='''{prompt}'''),
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
     retriever=dict(type=ZeroRetriever),
-    inferencer=dict(type=GenInferencer, max_out_len=512))
+    inferencer=dict(type=GenInferencer))
 
-cdme_eval_cfg = dict(
-    evaluator=dict(type=CDMEEvaluator),
-    pred_postprocessor=dict(type=cdme_postprocess),
-    dataset_postprocessor=dict(type=cdme_dataset_postprocess),
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchParallelEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
     pred_role='BOT')
 
-context_lengths = list(range(1000, 33000, 1000))
+context_lengths = list(range(1000, 5000, 1000))
 document_depth_percent_intervals = 20
 document_depth_percent_interval_type = "linear"
 
-base_path = './data/CDME'
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
+depths_float = generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type)
+depths = [int(depth) for depth in depths_float]
+
+for original_context_length in context_lengths:
+    dataset_dict = {
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_en_4k',
+        'type': NeedleBenchParallelDataset,
+        'path': base_path,
+        'needle_file_name': needle_file_name,
+        'length': original_context_length,
+        'depths': depths,
+        'tokenizer_model': 'gpt-4',
+        'file_list': file_list,
+        'num_repeats_per_file': 25,
+        'length_buffer': 1000,
+        'guide': True,
+        'language': 'English',
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
+    }
+    needlebench_datasets_en.append(dataset_dict)
+
 file_list = ['zh_finance.jsonl']
-cdme_datasets = []
+needlebench_datasets_zh = []
 
 for original_context_length in context_lengths:
-    for depth_percent in generate_depth_percents(
-            document_depth_percent_intervals,
-            document_depth_percent_interval_type):
     dataset_dict = {
-            'abbr': f'CDME_Length{original_context_length}'
-            f'Depth{int(depth_percent)}',
-            'type': CDMEDataset,
+        'abbr': f'Length{original_context_length}'
+                f'_parallel_zh_4k',
+        'type': NeedleBenchParallelDataset,
         'path': base_path,
+        'needle_file_name': needle_file_name,
         'length': original_context_length,
-            'depth': int(depth_percent),
+        'depths': depths,
         'tokenizer_model': 'gpt-4',
         'file_list': file_list,
-            'num_repeats_per_file': 10,
+        'num_repeats_per_file': 25,
         'length_buffer': 200,
         'guide': True,
         'language': 'Chinese',
-            'needle': '\n小明最喜欢的实习的地点就是上海人工智能实验室。\n',
-            'retrieval_question': '小明最喜欢的实习地点是哪里？请按照'
-            '“小明最喜欢的实习地点就是________。”的格式回答。',
-            'reader_cfg': cdme_reader_cfg,
-            'infer_cfg': cdme_infer_cfg,
-            'eval_cfg': cdme_eval_cfg
+        'reader_cfg': needlebench_reader_cfg,
+        'infer_cfg': needlebench_infer_cfg,
+        'eval_cfg': needlebench_eval_cfg
     }
-        cdme_datasets.append(dataset_dict)
+    needlebench_datasets_zh.append(dataset_dict)

configs/datasets/cdme/cdme200k.py → configs/datasets/needlebench/needlebench_4k/needlebench_single.py

 from opencompass.openicl.icl_prompt_template import PromptTemplate
 from opencompass.openicl.icl_retriever import ZeroRetriever
 from opencompass.openicl.icl_inferencer import GenInferencer
-from opencompass.datasets.cdme.cdme import CDMEDataset
-from opencompass.datasets.cdme.cdme import CDMEEvaluator
-from opencompass.datasets.cdme.cdme import cdme_postprocess
-from opencompass.datasets.cdme.cdme import cdme_dataset_postprocess
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginDataset
+from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
 import math
 
 
@@ -31,67 +31,72 @@ def generate_depth_percents(intervals, interval_type):
         raise ValueError('Unsupported interval type')
 
 
-cdme_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
 
-cdme_infer_cfg = dict(
+needlebench_infer_cfg = dict(
     prompt_template=dict(
         type=PromptTemplate,
-        template='''{prompt}'''),
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
     retriever=dict(type=ZeroRetriever),
-    inferencer=dict(type=GenInferencer, max_out_len=512))
+    inferencer=dict(type=GenInferencer))
 
-cdme_eval_cfg = dict(
-    evaluator=dict(type=CDMEEvaluator),
-    pred_postprocessor=dict(type=cdme_postprocess),
-    dataset_postprocessor=dict(type=cdme_dataset_postprocess),
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchOriginEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
     pred_role='BOT')
 
-cdme_trim_eval_cfg = dict(
-    evaluator=dict(type=CDMEEvaluator, use_trim=True),
-    pred_postprocessor=dict(type=cdme_postprocess),
-    dataset_postprocessor=dict(type=cdme_dataset_postprocess),
-    pred_role='BOT')
-
-#context_lengths = list(range(1000, 201000, 1000))
-context_lengths = [16000, 32000, 48000, 64000, 80000, 96000, 112000, 128000, 144000, 160000, 176000, 192000, 200000]
+context_lengths = list(range(1000, 5000, 1000))
 document_depth_percent_intervals = 20
 document_depth_percent_interval_type = "linear"
 
-base_path = './data/CDME'
-file_list = ['zh_finance.jsonl']
-cdme_datasets = []
-cdme_trim_datasets = []
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+needlebench_datasets_en = []
+needle_file_name = 'needles.jsonl'
 
 for original_context_length in context_lengths:
     for depth_percent in generate_depth_percents(
             document_depth_percent_intervals,
             document_depth_percent_interval_type):
         dataset_dict = {
-            'abbr': f'CDME_Length{original_context_length}'
-            f'Depth{int(depth_percent)}',
-            'type': CDMEDataset,
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_en_4k',
+            'type': NeedleBenchOriginDataset,
             'path': base_path,
             'length': original_context_length,
             'depth': int(depth_percent),
             'tokenizer_model': 'gpt-4',
             'file_list': file_list,
             'num_repeats_per_file': 10,
-            'length_buffer': 200,
+            'length_buffer': 600,
             'guide': True,
-            'language': 'Chinese',
-            'needle': '\n小明最喜欢的实习的地点就是上海人工智能实验室。\n',
-            'retrieval_question': '小明最喜欢的实习地点是哪里？请按照'
-            '“小明最喜欢的实习地点就是________。”的格式回答。\n',
-            'reader_cfg': cdme_reader_cfg,
-            'infer_cfg': cdme_infer_cfg,
-            'eval_cfg': cdme_eval_cfg
+            'language': 'English',
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
         }
-        cdme_datasets.append(dataset_dict)
+        needlebench_datasets_en.append(dataset_dict)
+
+file_list = ['zh_finance.jsonl']
+needlebench_datasets_zh = []
+needle_file_name = 'needles.jsonl'
 
-        trim_dataset_dict = {
-            'abbr': f'CDME_Length{original_context_length}'
-            f'Depth{int(depth_percent)}',
-            'type': CDMEDataset,
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_origin_zh_4k',
+            'type': NeedleBenchOriginDataset,
             'path': base_path,
             'length': original_context_length,
             'depth': int(depth_percent),
@@ -101,11 +106,9 @@ for original_context_length in context_lengths:
             'length_buffer': 200,
             'guide': True,
             'language': 'Chinese',
-            'needle': '\n小明最喜欢的实习的地点就是上海人工智能实验室。\n',
-            'retrieval_question': '小明最喜欢的实习地点是哪里？请按照'
-            '“小明最喜欢的实习地点就是________。”的格式回答。\n',
-            'reader_cfg': cdme_reader_cfg,
-            'infer_cfg': cdme_infer_cfg,
-            'eval_cfg': cdme_trim_eval_cfg
+            'needle_file_name': needle_file_name,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
         }
-        cdme_trim_datasets.append(trim_dataset_dict)
+        needlebench_datasets_zh.append(dataset_dict)

configs/datasets/needlebench/needlebench_8k/needlebench.py

+from mmengine.config import read_base
+
+with read_base():
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_en as needlebench_multi_2needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_en as needlebench_multi_3needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_en as needlebench_multi_4needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_en as needlebench_multi_5needle_en_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_2needle_zh as needlebench_multi_2needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_3needle_zh as needlebench_multi_3needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_4needle_zh as needlebench_multi_4needle_zh_datasets
+    from .needlebench_multi_reasoning import needlebench_datasets_5needle_zh as needlebench_multi_5needle_zh_datasets
+
+    from .needlebench_single import needlebench_datasets_en as needlebench_origin_en_datasets
+    from .needlebench_single import needlebench_datasets_zh as needlebench_origin_zh_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_en as needlebench_parallel_en_datasets
+    from .needlebench_multi_retrieval import needlebench_datasets_zh as needlebench_parallel_zh_datasets
+
+needlebench_datasets = sum((v for k, v in locals().items() if k.endswith('_datasets')), [])

configs/datasets/needlebench/needlebench_8k/needlebench_multi_reasoning.py

+from opencompass.openicl.icl_prompt_template import PromptTemplate
+from opencompass.openicl.icl_retriever import ZeroRetriever
+from opencompass.openicl.icl_inferencer import GenInferencer
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiDataset
+from opencompass.datasets.needlebench.multi import NeedleBenchMultiEvaluator
+from opencompass.datasets.needlebench.origin import needlebench_postprocess
+from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess
+import math
+
+
+def logistic(x, L=100, x0=50, k=0.1):
+    return round(L / (1 + math.exp(-k * (x - x0))), 3)
+
+
+def generate_linear_space(start, end, num):
+    if num == 1:
+        return [start]
+    elif num < 1:
+        raise ValueError("num must be at least 1.")
+    step = (end - start) / (num - 1)
+    return [start + step * i for i in range(num)]
+
+
+def generate_depth_percents(intervals, interval_type):
+    if interval_type == 'linear':
+        return generate_linear_space(0, 100, intervals)
+    elif interval_type == 'sigmoid':
+        linear_space = generate_linear_space(0, 100, intervals)
+        return [logistic(x) for x in linear_space]
+    else:
+        raise ValueError('Unsupported interval type')
+
+
+needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')
+
+needlebench_infer_cfg = dict(
+    prompt_template=dict(
+        type=PromptTemplate,
+        template=dict(
+            round=[
+                dict(role='HUMAN', prompt='{prompt}'),
+                dict(role='BOT', prompt='{answer}\n'),
+            ]
+        )
+        ),
+    retriever=dict(type=ZeroRetriever),
+    inferencer=dict(type=GenInferencer))
+
+needlebench_eval_cfg = dict(
+    evaluator=dict(type=NeedleBenchMultiEvaluator),
+    pred_postprocessor=dict(type=needlebench_postprocess),
+    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
+    pred_role='BOT')
+
+context_lengths = list(range(5000, 9000, 1000))
+document_depth_percent_intervals = 20
+document_depth_percent_interval_type = "linear"
+
+# ----------English Version----------
+base_path = './data/needlebench'
+file_list = ['PaulGrahamEssays.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_en.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_en = []
+language = 'English'
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 1000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_en.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 1000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_en.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 1000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_en.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_en = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_en_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 1000,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_en.append(dataset_dict)
+
+# ----------Chinese Version----------
+base_path = './data/needlebench'
+file_list = ['zh_finance.jsonl']
+
+needle_file_name = 'multi_needle_reasoning_zh.json'
+diff = 10
+num_needles = 2
+needlebench_datasets_2needle_zh = []
+language = 'Chinese'
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_2needle_zh.append(dataset_dict)
+
+num_needles = 3
+needlebench_datasets_3needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_3needle_zh.append(dataset_dict)
+
+num_needles = 4
+needlebench_datasets_4needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_4needle_zh.append(dataset_dict)
+
+num_needles = 5
+needlebench_datasets_5needle_zh = []
+
+for original_context_length in context_lengths:
+    for depth_percent in generate_depth_percents(
+            document_depth_percent_intervals,
+            document_depth_percent_interval_type):
+        dataset_dict = {
+            'abbr': f'Length{original_context_length}'
+                    f'Depth{int(depth_percent)}_{num_needles}needle_zh_8k',
+            'type': NeedleBenchMultiDataset,
+            'path': base_path,
+            'length': original_context_length,
+            'depth': int(depth_percent),
+            'tokenizer_model': 'gpt-4',
+            'file_list': file_list,
+            'num_repeats_per_file': 10,
+            'length_buffer': 200,
+            'guide': True,
+            'language': language,
+            'needle_file_name': needle_file_name,
+            'num_needles': num_needles,
+            'diff': diff,
+            'reader_cfg': needlebench_reader_cfg,
+            'infer_cfg': needlebench_infer_cfg,
+            'eval_cfg': needlebench_eval_cfg
+        }
+        needlebench_datasets_5needle_zh.append(dataset_dict)