# Inspect Task Evaluation README

## Overview

The Inspect implementation for NIAH (Needle in a Haystack) is designed to evaluate the in-context retrieval ability of long context language models (LLMs). The main goal of this evaluation is to place a random fact or statement (the "needle") in the middle of a long context window (the "haystack") and ask the model to retrieve this statement. The evaluation iterates over various document depths and context lengths to measure performance.

This evaluation was contributed by [Owen Parsons](https://owenparsons.github.io).

## Installation

To get started, install the required Python packages:

```bash
pip install inspect_ai
pip install git+https://github.com/UKGovernmentBEIS/inspect_evals
```

## Running the Evaluation

You can evaluate against models using the following command:

```bash
inspect eval inspect_evals/niah --model openai/gpt-3.5-turbo-16k
```

If you prefer not to specify the `--model` argument each time you run an evaluation, create a `.env` configuration file in your working directory that defines the `INSPECT_EVAL_MODEL` environment variable along with your API key.

### Example of the .env file:
```
INSPECT_EVAL_MODEL=openai/gpt-3.5-turbo-16k
```

You can also use the `-T` tag to define input values for the evaluation. For example:

```bash
inspect eval niah.py --model openai/gpt-3.5-turbo -T n_needles=20
```

## Configuration Variables

Here are the configuration variables used in the NIAH evaluation, along with their default values and types:

| Variable          | Type     | Default Value | Description                                                                                   |
|-------------------|----------|---------------|-----------------------------------------------------------------------------------------------|
| `min_context`     | `int`    | `10000`         | Minimum context length to evaluate.                                                          |
| `max_context`     | `int`    | `120000`         | Maximum context length to evaluate.                                                          |
| `n_contexts`      | `int`    | `15`           | The number of contexts to evaluate.                                                          |
| `n_positions`      | `int`    | `15`           | The number of positions to evaluate for a given context length.                             |
| `start_buffer`      | `int`    | `0`           | Buffer at the top of the context to avoid placing needles. (Example: If `start_buffer` is `100`, then the first needle position would aim to be at the 100th token in the context.)                                   |
| `end_buffer`   | `int`    | `0`           | Buffer at the bottom of the context to avoid placing needles. (Example: For a context length of `1000`, if `end_buffer` is `100`, then the final needle position would aim to be at the 900th token in the context.)                                |
| `n_needles`       | `int`    | `1`          | The number of needles to sample.                                                             |
| `sample_method`   | `str`    | `"fixed"`    | Method for sampling the needles.                                                              |
| `fixed_index`     | `int`    | `0`           | The index of the needle to use when `sample_method` is set to `"fixed"`.                    |
| `n_runs`          | `int`    | `1`           | The number of runs for the evaluation.                                                       |

## Scoring Metric

This benchmark uses a modified version of ``model_graded_qa()`` that allows for the scorer call to only include the question related to the needle, rather than passing the original prompt during the benchmark task. This is to avoid the model having to handle long context inputs during the scoring.

The scoring criteria is taken from [Greg Kramradt's implementation](https://github.com/gkamradt/LLMTest_NeedleInAHaystack) and is shown below:

```
Score 1: The answer is completely unrelated to the reference.
Score 3: The answer has minor relevance but does not align with the reference.
Score 5: The answer has moderate relevance but contains inaccuracies.
Score 7: The answer aligns with the reference but has minor omissions.
Score 10: The answer is completely accurate and aligns perfectly with the reference.
```

## Dataset

The dataset used for the evaluation is read from [OpenCompass](https://opencompass.readthedocs.io/en/latest/advanced_guides/needleinahaystack_eval.html), and the specific dataset is generated based on the values defined in the configuration.

## Dataset Construction

The final dataset used by the NIAH evaluation is generated using the OpenCompass dataset and the configuration variables by following the series of steps summarised below.

1. Model Encoding: An encoder for the specified model is created using the tiktoken library, which facilitates tokenisation for text processing.

2. Context Length Generation: Context lengths are generated for the specified range of context values given by `min_context`, `max_context` and `n_contexts`.

3. Needle Position Generation: Needle positions are determined across the generated contexts based on the number of needles specified by `n_needles` accounting for a buffer at the start/end of the document if specified by
`start_buffer` or `end_buffer`.

4. Data Retrieval: The relevant haystack and needle datasets are extracted from the Hugging Face dataset repository.

5. Needle Filtering: The needles are filtered to include only English entries.

6. Needle Sampling: Needles are sampled based on the chosen method (fixed, sequential, or random).

7. Needle Adjustment: The needles are repeated and shifted to prepare for multiple runs.

8. Token Length Calculations: The maximum combined token lengths for the needles, questions, and answers are computed, along with the token counts for the main and question prompts.

9. Context Reading and Trimming: The relevant context texts are read, combined and trimmed to match the required context token lengths.

10. Period Indices Calculation: Period indices are determined for each context length to facilitate the insertion of needles at appropriate positions between sentences.

11. Model Context Length Verification: A check is performed to ensure that the model context length is larger than the specified context lengths.

12. Dataset Creation: The final dataset is constructed by combining the contexts, needle insertions, and associated metadata.

13. DataFrame Conversion: The complete dataset is converted into a Pandas DataFrame.

## Experimental Results

### Procedure
Results were obtained for the NIAH benchmark using the `mistral-small-latest` and `mistral-medium-latest` models. Scores were obtained for multiple runs and averaged to obtain the final scores presented in the plots below (*n_runs* = 10 for `mistral-small-latest`, *n_runs* = 5 for `mistral-medium-latest`). For results for both models, 10 different insertion positions were used from 0% to 100%. This refers to the relative position of the inserted fact ('needle') with respect to the main context text ('haystack'). A total of 9 different context lengths were used for each model, ranging between 2000 and 27000 tokens.

### Model Calls
In the NIAH evaluation, there are two distinct model calls: one when the full context is passed to the model and the model is asked to retrieve the needle and one when the target answer and actual answer are passed to the model and the model is asked to evaluate the answer. I will refer to the two calls to the language model using different terms:
- **Task model**: This refers to the initial *retrieval* call to the model where the full context is passed and the model is asked to retrieve the needle.
- **Evaluation model**: This refers to the *scoring* call to the model where the target answer and actual answer are passed to the model and the model is required to grade the answer from 0 to 10. Note that the evaluation model is not required to be the same as the task model but for the results presented below, the same model is used for both task and evaluation.

### Small Model Results

The mean and standard deviation heatmaps for the `mistral-small-latest` model are shown in figure 1 below. Performance was generally good across most context lengths and needle positions, with a slight drop in performance seen in the largest context lengths. The experimental parameter pairs where performance was poor had relatively high variance across runs, indicating that the model was more sensitive to variations in the needle or haystack content.

<div style="display: flex; justify-content: space-between; gap: 20px;">
    <figure style="margin: 0; flex: 1;">
        <img src="images/mistral-small-latest_mean_heatmap.png" alt="Small Model Heatmap" width="100%">
    </figure>
    <figure style="margin: 0; flex: 1;">
        <img src="images/mistral-small-latest_std_heatmap.png" alt="Medium Model Heatmap" width="100%">
    </figure>
</div>
<figcaption align="center"><em style="font-size: 0.8em;">Figure 1: Mean (left) and standard deviation (right) heatmaps showing the performance of the mistral-small-latest model across different context lengths and needle positions.</em></figcaption>

### Medium Model Results

The mean and standard deviation heatmaps for the `mistral-medium-latest` model are shown in figure 2 below. Relative to the small model, the medium model actually has a drop in performance across most context lengths and needle positions with mean scores tending to be lower. This was unexpected behaviour, given that the medium model should have a deeper understanding of the input context and would be expected to be able to retrieve the needle more accurately. The higher standard deviation values suggest that there is a greater variation in performance across runs, indicating that the medium model might be particularly sensitive to variations in the needle or haystack content. Notably, the performance of the medium model was also found to be sensitive to the insertion position, with the model showing particularly poor performance when the needled was inserted at the end of the context window.

<div style="display: flex; justify-content: space-between; gap: 20px;">
    <figure style="margin: 0; flex: 1;">
        <img src="images/mistral-medium-latest_mean_heatmap.png" alt="Small Model Heatmap" width="100%">
    </figure>
    <figure style="margin: 0; flex: 1;">
        <img src="images/mistral-medium-latest_std_heatmap.png" alt="Medium Model Heatmap" width="100%">
    </figure>
</div>
<figcaption align="center"><em style="font-size: 0.8em;">Figure 2: Mean (left) and standard deviation (right) heatmaps showing the performance of the mistral-medium-latest model across different context lengths and needle positions.</em></figcaption>

### Comparative Results

The [original implementation of the NIAH benchmark](https://github.com/gkamradt/LLMTest_NeedleInAHaystack) presented results for OpenAI's GPT-4-128k model and Anthropic's Claude 2.1 model. Notably, the results for the Claude 2.1 model were surprising, with the model tending to score much worse than GPT-4-128k. This was followed up [in a post by Anthropic](https://www.anthropic.com/news/claude-2-1-prompting), which demonstrated how a small change to the prompt used to generate the haystack could lead to a drastic improvement in performance. 

[Additional results were posted by Arize AI](https://arize.com/blog-course/the-needle-in-a-haystack-test-evaluating-the-performance-of-llm-rag-systems/) for the GPT-4-128k model and two Mistral models (Mixtral 8X7B-v0.1 and 7B Instruct). While the results for the Mistral models are not directly comparable to the results presented above, there is an indication that performance drops for positions very close to the end of the context window similarly to the results from the `mistral-medium-latest` model.


### Exploring Responses from the Evaluation Model

To probe a bit deeper into the *evaluation model* responses, the distribution of scores given by the evaluation model were plotted (see figure 3 below). Results are shown for both the `mistral-small-latest` and `mistral-medium-latest` models. The distribution shows that responses from the *evaluation model* tended to either be the highest (10) or lowest scores (1), with a few responses in the middle (Note that the mistral-small-latest model gave some responses of 0, even though the prompt specified a range of 1-10). These results suggest that the *evaluation model* in the NIAH evaluation tends to operate in a fairly binary manner, with responses tending to be either completely correct or completely incorrect.

<figure>
    <img src="images/combined_models_score_distribution.png" alt="Distribution of scores" width="600">
    <figcaption align="center"><em style="font-size: 0.8em;">Figure 3: Distribution of scores given by the evaluation model shown for the `mistral-small-latest` and `mistral-medium-latest` models. Frequency (y-axis) is on a log scale with absolute counts shown above the bars.</em></figcaption>
</figure>

To further explore the *evaluation model* responses, manual inspection of responses was carried out. After manually examining a number of the *evaluation model* responses, it was found that there were numerous responses where the *evaluation model* was scoring correct answers as incorrect/irrelevant (by assigning lower scores) due to the fact that the content of the answer seemed fictional or was not in line with the models understanding of the world. One such example is shown below:

```
"The submission is completely unrelated to the criterion as it suggests 'the Virtual Rocket Band' performed on the Moon, which is inaccurate. No band has performed on the Moon to date. The question may have been intended as a hypothetical or fictional scenario, but without that context, the answer does not align with reality."
```

This represents a failure of the *evaluation model* and may be a result of the prompt not stating clearly enough that the only the target answer should be used and prior knowledge should not be factored into the evaluation. This occurrence was more common in the `mistral-medium-latest` model relative to the `mistral-small-latest` model and might explain why the results indicate poorer performance on the benchmark in the `mistral-medium-latest` model relative to the `mistral-small-latest` model.

To understand how much this type of response might be impacting the results, I extracted the explanations from the *evaluation model* for all queries and flagged the ones where this type of 'unrelated answer' response occurred. This was done simply by searching for the string 'unrelated' in the explanation field of the *evaluation model* responses and was carried out for the `mistral-medium-latest` model. The results indicated that between 10% and 20% of the responses were flagged as 'unrelated' for most of the insertion positions, with a spike in the number of 'unrelated' responses seen when the insertion position was at the end of the main context. The results are shown in the plot below (figure 4):

<figure>
    <img src="images/mistral-medium-latest_unrelated_answer_responses.png" alt="Unrelated Responses Plot" width="600">
    <figcaption align="center"><em style="font-size: 0.8em;">Figure 4: The percentage of responses that were flagged as 'unrelated' for the `mistral-medium-latest` model as a function of the target position. Results were averaged across all context lengths and multiple runs. The dashed line indicates the average score for the model and the shaded region represents the standard deviation.</em></figcaption>
</figure>

Note, this is by no means a straightforward or intuitive result. The insertion position should have no direct effect on the *evaluation model* as it only has access to the answers and target solutions and does not have direct access to the context. Therefore, any effect observed here must be due to differences in the answers provided by the initial *task model* for a given position. Additional manual inspection of responses from the *task model* revealed that the answers provided when the insertion position was at the end of the main context were more likely to be suffixed with clarifying statements that indicated the information used to answer the question was unrelated to the rest of the context. Some examples are show below:

```
The first band to perform on the Moon was the Virtual Rocket Band.


Note: This information is not found in the provided context. It seems to be an unrelated statement at the end of the text.
```

```
The first band to perform on the Moon was the Virtual Rocket Band.

Note: This answer is a joke, as there is no mention of any band performing on the Moon in the provided context. The last sentence appears to be a non-sequitur and is not related to the previous discussion.
```

I suspect this might be due to the model being trained in a way that guides it away from immediately repeating the most recent text from the input prompt in it's reply. These results highlight the importance of carefully evaluating the results from the NIAH benchmark rather than relying on summarised or aggregated scores. Different models will likely have different 'weak spots' in the `position × context length` parameter space based on nuances that come from differences in training procedures or system prompts.
