jonathan-roberts1/needle-threading
Needle Threading: Can LLMs Follow Threads through Near-Million-Scale Haystacks? Dataset Summary As the context limits of Large Language Models (LLMs) increase, the range of possible applications and downstream functions broadens. Although the development of longer context models has seen rapid gains recently, our understanding of how effectively they use their context has not kept pace. To address this, we conduct a set of retrieval experiments designed to… See the full description on the dataset page: https://huggingface.co/datasets/jonathan-roberts1/needle-threading.
Needle Threading: Can LLMs Follow Threads through Near-Million-Scale Haystacks?
Dataset Description
- Homepage: https://needle-threading.github.io
- Paper: Needle Threading: Can LLMs Follow Threads through Near-Million-Scale Haystacks?
- Repository Needle Threading
Dataset Summary
As the context limits of Large Language Models (LLMs) increase, the range of possible applications and downstream functions broadens. Although the development of longer context models has seen rapid gains recently, our understanding of how effectively they use their context has not kept pace.
To address this, we conduct a set of retrieval experiments designed to evaluate the capabilities of 17 leading LLMs, such as their ability to follow threads of information through the context window.
Strikingly, we find that many models are remarkably thread-safe: capable of simultaneously following multiple threads without significant loss in performance. Still, for many models, we find the effective context limit is significantly shorter than the supported context length, with accuracy decreasing as the context window grows.
Example Usage
Option 1: HuggingFace datasets
from datasets import load_dataset
# task splits can be downloaded separately:
# splits = ['Single_Needle', 'Multi_Needle', 'Conditional_Needle', 'Single_Thread', 'Multi_Thread']
single_needle_dataset = load_dataset("jonathan-roberts1/needle-threading", split='Single_Needle')
"""
Dataset({
features: ['id', 'haystack', 'keys', 'values', 'question', 'context_length', 'num_kv_pairs',
'repeat_number', 'needle_depth', 'num_needles', 'needle_placement', 'conditional_character',
'thread_length', 'thread_direction', 'num_threads'],
num_rows: 660
})
Note the units of context_length are number of characters.
"""
# query individual questions
single_needle_dataset[5] # e.g., the 6th element
"""
{'id': 5, 'haystack': '{"e3e70682-c209-4cac-629f-6fbed82c07cd": "f728b4fa-4248-5e3a-0a5d-2f346baa9455",
"eb1...": "964a870c-7c87-9b74-1d87-8f9f9cdf5a86"}', 'keys': '247a8333-f7b0-b7d2-cda8-056c3d15eef7',
'values': '1759edc3-72ae-2244-8b01-63c1cd9d2b7d', 'question': 'Extract the value corresponding to
the specified key in the JSON object. Key: "247a83...-cda8-056c3d15eef7"\n Corresponding value: ',
'context_length': 2000, 'num_kv_pairs': 25, 'repeat_number': 0, 'needle_depth': '50', 'num_needles': 1,
'needle_placement': 'depthwise', 'conditional_character': 'N/A', 'thread_length': 1,
'thread_direction': 'N/A', 'num_threads': 0}
"""Option 2: Manual download
Directly downloading image files and question data from the needle-threading HuggingFace repository into the ``data`` directory in this repo.
cd data
wget https://huggingface.co/datasets/jonathan-roberts1/needle-threading/resolve/main/data_json.zip
unzip data_json.zip && rm data_json.zipExpected structure
├── data
├── json_data
├── Single_Needle.json
├── Multiple_Needles.json
├── Conditional_Needles.json
├── Single_Threads.json
├── Multi_Threads.jsonNote: ``data_json/`` needs to be downloaded.
Please visit our GitHub repository for example inference code.
Dataset Curators
This dataset was curated by Jonathan Roberts, Kai Han, and Samuel Albanie
Citation
If you found our work useful in your own research, please consider citing our paper:
@article{roberts2024needle,
title={Needle Threading: Can LLMs Follow Threads through Near-Million-Scale Haystacks?},
author={Roberts, Jonathan and Han, Kai and Albanie, Samuel},
journal={arXiv preprint arXiv:2411.05000},
year={2024}
}