"Multi-Needle In A Haystack" でLLMのロングコンテキストの評価を試す

!pip install needlehaystack

import os
from google.colab import userdata

os.environ["NIAH_EVALUATOR_API_KEY"] = userdata.get('OPENAI_API_KEY')

os.environ["NIAH_MODEL_API_KEY"] = userdata.get('OPENAI_API_KEY')

!needlehaystack.run_test --help

usage: needlehaystack.run_test [-h] [--config CONFIG] [--print_config[=flags]]
                               [--provider PROVIDER] [--evaluator EVALUATOR]
                               [--model_name MODEL_NAME]
                               [--evaluator_model_name EVALUATOR_MODEL_NAME] [--needle NEEDLE]
                               [--haystack_dir HAYSTACK_DIR]
                               [--retrieval_question RETRIEVAL_QUESTION]
                               [--results_version RESULTS_VERSION]
                               [--context_lengths_min CONTEXT_LENGTHS_MIN]
                               [--context_lengths_max CONTEXT_LENGTHS_MAX]
                               [--context_lengths_num_intervals CONTEXT_LENGTHS_NUM_INTERVALS]
                               [--context_lengths CONTEXT_LENGTHS]
                               [--document_depth_percent_min DOCUMENT_DEPTH_PERCENT_MIN]
                               [--document_depth_percent_max DOCUMENT_DEPTH_PERCENT_MAX]
                               [--document_depth_percent_intervals DOCUMENT_DEPTH_PERCENT_INTERVALS]
                               [--document_depth_percents DOCUMENT_DEPTH_PERCENTS]
                               [--document_depth_percent_interval_type DOCUMENT_DEPTH_PERCENT_INTERVAL_TYPE]
                               [--num_concurrent_requests NUM_CONCURRENT_REQUESTS]
                               [--save_results {true,false,null}]
                               [--save_contexts {true,false,null}]
                               [--final_context_length_buffer FINAL_CONTEXT_LENGTH_BUFFER]
                               [--seconds_to_sleep_between_completions SECONDS_TO_SLEEP_BETWEEN_COMPLETIONS]
                               [--print_ongoing_status {true,false,null}] [--eval_set EVAL_SET]
                               [--multi_needle {true,false,null}] [--needles NEEDLES]

<class 'needlehaystack.run.CommandArgs'>

options:
  -h, --help            Show this help message and exit.
  --config CONFIG       Path to a configuration file.
  --print_config[=flags]
                        Print the configuration after applying all other arguments and exit. The
                        optional flags customizes the output and are one or more keywords
                        separated by comma. The supported flags are: comments, skip_default,
                        skip_null.
  --provider PROVIDER   (type: str, default: openai)
  --evaluator EVALUATOR
                        (type: str, default: openai)
  --model_name MODEL_NAME
                        (type: str, default: gpt-3.5-turbo-0125)
  --evaluator_model_name EVALUATOR_MODEL_NAME
                        (type: Optional[str], default: gpt-3.5-turbo-0125)
  --needle NEEDLE       (type: Optional[str], default: The best thing to do in San Francisco is
                        eat a sandwich and sit in Dolores Park on a sunny day. )
  --haystack_dir HAYSTACK_DIR
                        (type: Optional[str], default: PaulGrahamEssays)
  --retrieval_question RETRIEVAL_QUESTION
                        (type: Optional[str], default: What is the best thing to do in San
                        Francisco?)
  --results_version RESULTS_VERSION
                        (type: Optional[int], default: 1)
  --context_lengths_min CONTEXT_LENGTHS_MIN
                        (type: Optional[int], default: 1000)
  --context_lengths_max CONTEXT_LENGTHS_MAX
                        (type: Optional[int], default: 16000)
  --context_lengths_num_intervals CONTEXT_LENGTHS_NUM_INTERVALS
                        (type: Optional[int], default: 35)
  --context_lengths CONTEXT_LENGTHS, --context_lengths+ CONTEXT_LENGTHS
                        (type: Optional[list[int]], default: null)
  --document_depth_percent_min DOCUMENT_DEPTH_PERCENT_MIN
                        (type: Optional[int], default: 0)
  --document_depth_percent_max DOCUMENT_DEPTH_PERCENT_MAX
                        (type: Optional[int], default: 100)
  --document_depth_percent_intervals DOCUMENT_DEPTH_PERCENT_INTERVALS
                        (type: Optional[int], default: 35)
  --document_depth_percents DOCUMENT_DEPTH_PERCENTS, --document_depth_percents+ DOCUMENT_DEPTH_PERCENTS
                        (type: Optional[list[int]], default: null)
  --document_depth_percent_interval_type DOCUMENT_DEPTH_PERCENT_INTERVAL_TYPE
                        (type: Optional[str], default: linear)
  --num_concurrent_requests NUM_CONCURRENT_REQUESTS
                        (type: Optional[int], default: 1)
  --save_results {true,false,null}
                        (type: Optional[bool], default: True)
  --save_contexts {true,false,null}
                        (type: Optional[bool], default: True)
  --final_context_length_buffer FINAL_CONTEXT_LENGTH_BUFFER
                        (type: Optional[int], default: 200)
  --seconds_to_sleep_between_completions SECONDS_TO_SLEEP_BETWEEN_COMPLETIONS
                        (type: Optional[float], default: null)
  --print_ongoing_status {true,false,null}
                        (type: Optional[bool], default: True)
  --eval_set EVAL_SET   (type: Optional[str], default: multi-needle-eval-pizza-3)
  --multi_needle {true,false,null}
                        (type: Optional[bool], default: False)
  --needles NEEDLES, --needles+ NEEDLES
                        (type: list[str], default: [' Figs are one of the secret ingredients
                        needed to build the perfect pizza. ', ' Prosciutto is one of the secret
                        ingredients needed to build the perfect pizza. ', ' Goat cheese is one of
                        the secret ingredients needed to build the perfect pizza. '])

!needlehaystack.run_test \
  --provider openai \
  --evaluator openai \
  --model_name "gpt-3.5-turbo-0125" \
  --document_depth_percents "[50]" \
  --context_lengths "[2000]"

Testing single-needle


Starting Needle In A Haystack Testing...
- Model: gpt-3.5-turbo-0125
- Context Lengths: 1, Min: 2000, Max: 2000
- Document Depths: 1, Min: 50%, Max: 50%
- Needle: The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.



-- Test Summary -- 
Duration: 1.0 seconds
Context: 2000 tokens
Depth: 50%
Score: 10
Response: The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.

$ tree contexts
 contexts
└── gpt-3_5-turbo-0125_len_2000_depth_5000_context.txt

$ tree results
results
└── gpt-3_5-turbo-0125_len_2000_depth_5000_results.json

import tiktoken

def display_file_lines(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()
        
        encoding = tiktoken.get_encoding("cl100k_base")
        
        total_char_count = 0
        total_token_count = 0
        
        for i, line in enumerate(lines, start=1):
            line = line.strip()
            char_count = len(line)
            token_count = len(encoding.encode(line))
            
            total_char_count += char_count
            total_token_count += token_count
            
            print("{:4}:{:6}:{:6}: {}".format(
                i, total_char_count, total_token_count, line))


file_path = 'contexts/gpt-3_5-turbo-0125_len_2000_depth_5000_context.txt'
display_file_lines(file_path)

   1:   123:    26: May 2006(This essay is derived from a keynote at Xtech.)Could you reproduce Silicon Valley elsewhere, or is there something
   2:   202:    43: unique about it?It wouldn't be surprising if it were hard to reproduce in other
   3:   264:    56: countries, because you couldn't reproduce it in most of the US
   4:   392:    87: either.  What does it take to make a silicon valley even here?What it takes is the right people.  If you could get the right ten
   5:   455:    99: thousand people to move from Silicon Valley to Buffalo, Buffalo
(snip)
  48:  3550:   792: The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.
(snip)
 103:  7472:  1650: thesis seems to be that as wealth derives increasingly from ideas,
 104:  7539:  1664: cities will prosper only if they attract those who have them.  That
 105:  7608:  1678: is certainly true; in fact it was the basis of Amsterdam's prosperity
 106:  7689:  1695: 400 years ago.A lot of nerd tastes they share with the creative class in general.
 107:  7744:  1705: For example, they like well-preserved old neighborhoods

import json
from pprint import pprint 

with open("results/gpt-3_5-turbo-0125_len_2000_depth_5000_results.json") as f:
    pprint(json.load(f))

{
    'context_length': 2000,
    'depth_percent': 50.0,
    'file_name': 'gpt-3_5-turbo-0125_len_2000_depth_5000',
    'model': 'gpt-3.5-turbo-0125',
    'model_response': 'The best thing to do in San Francisco is eat a sandwich '
                      'and sit in Dolores Park on a sunny day.',
    'needle': '\n'
              'The best thing to do in San Francisco is eat a sandwich and sit in '
              'Dolores Park on a sunny day.\n',
    'score': 10,
    'test_duration_seconds': 0.9612917900085449,
    'test_timestamp_utc': '2024-03-17 15:16:21+0000',
    'version': 1
}

os.environ["NIAH_MODEL_API_KEY"] = userdata.get('ANTHROPIC_API_KEY')

!needlehaystack.run_test \
    --provider anthropic \
    --evaluator openai \
    --model_name "claude-2.1" \
    --document_depth_percents "[50]" \
    --context_lengths "[2000]"

Testing single-needle


Starting Needle In A Haystack Testing...
- Model: claude-2.1
- Context Lengths: 1, Min: 2000, Max: 2000
- Document Depths: 1, Min: 50%, Max: 50%
- Needle: The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.



-- Test Summary -- 
Duration: 2.9 seconds
Context: 2000 tokens
Depth: 50%
Score: 10
Response: 

"The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day."

$ tree contexts
contexts/
├── claude-2_1_len_2000_depth_5000_context.txt
└── gpt-3_5-turbo-0125_len_2000_depth_5000_context.txt

$ tree results
results
├── claude-2_1_len_2000_depth_5000_results.json
└── gpt-3_5-turbo-0125_len_2000_depth_5000_results.json

$ diff claude-2_1_len_2000_depth_5000_context.txt gpt-3_5-turbo-0125_len_2000_depth_5000_context.txt
47c47,49
< them to make their own investment decisions.  Most are only allowed
---
> them to make their own investment decisions.
> The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.
>   Most are only allowed
50,51d51
< The best thing to do in San Francisco is eat a sandwich and sit in Dolores Park on a sunny day.
< 
104c104,107
< cities will prosper only if they attract those
\ No newline at end of file
---
> cities will prosper only if they attract those who have them.  That
> is certainly true; in fact it was the basis of Amsterdam's prosperity
> 400 years ago.A lot of nerd tastes they share with the creative class in general.
> For example, they like well-preserved old neighborhoods
\ No newline at end of file

%%bash

context_length=1000
max_context_length=120000
step=8000

while [ $context_length -le $max_context_length ]; do
    echo needlehaystack.run_test \
    --provider anthropic \
    --evaluator openai \
    --model_name "claude-2.1" \
    --document_depth_percents "[0,10,20,30,40,50,60,70,80,90,100]" \
    --context_lengths "[${context_length}]"
    next_exp_num=$((context_length * 2))
    if [ $((next_exp_num - context_length)) -le $step ]; then
        context_length=$next_exp_num
    else
        context_length=$((context_length + step))
    fi
done

needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [1000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [2000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [4000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [8000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [16000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [24000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [32000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [40000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [48000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [56000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [64000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [72000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [80000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [88000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [96000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [104000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [112000]
needlehaystack.run_test --provider anthropic --evaluator openai --model_name claude-2.1 --document_depth_percents [0,10,20,30,40,50,60,70,80,90,100] --context_lengths [120000]

!needlehaystack.run_test \
  --provider openai \
  --evaluator openai \
  --model_name "gpt-3.5-turbo-0125" \
  --haystack_dir "./contexts_jp" \     # 日本語コンテキストのディレクトリ
  --needle "サンフランシスコでの一番の楽しみは、晴れた日にドローレス公園でサンドイッチを食べることです。" \
  --retrieval_question "サンフランシスコでの一番の楽しみは？" \
  --document_depth_percents "[50]" \
  --context_lengths "[2000]"