---
title: "in-context-ralm"
type: "tool"
slug: "ai21labs-in-context-ralm"
canonical_url: "https://www.graphcanon.com/tools/ai21labs-in-context-ralm"
github_url: "https://github.com/AI21Labs/in-context-ralm"
homepage_url: null
stars: 295
forks: 28
primary_language: "Python"
license: "Apache-2.0"
archived: true
categories: ["model-training", "evaluation-observability"]
tags: ["wikitext-103", "question-answering-experiments", "language-models", "retrieval-augmentation", "transformers", "pyserini", "bm25"]
updated_at: "2026-07-12T00:04:54.251936+00:00"
---

# in-context-ralm

> In-Context Retrieval-Augmented Language Models

> **Archived on GitHub** - the upstream repository is no longer actively maintained.

Repository for reproducing experiments on WikiText-103 from AI21 Labs' research paper, focusing on in-context retrieval-augmented language models.

## Facts

- Repository: https://github.com/AI21Labs/in-context-ralm
- Stars: 295 · Forks: 28 · Open issues: 4 · Watchers: 6
- Primary language: Python
- License: Apache-2.0
- Last pushed: 2023-12-20T07:55:54+00:00

## Trust & health

_Signals computed from public GitHub metadata. Not a security guarantee._

- Maintenance: Archived (computed 2026-07-11T23:06:10.486Z)
- Security scan: Findings present (0 critical, 0 high, 0 medium, 75 low) · last scan 2026-07-11T23:06:10.960Z
- Full report: [trust report](/tools/ai21labs-in-context-ralm/trust.md) · [JSON](https://www.graphcanon.com/api/graphcanon/tools/ai21labs-in-context-ralm/trust)

## Categories

- [Model Training](/categories/model-training.md)
- [Evaluation & Observability](/categories/evaluation-observability.md)

## Tags

wikitext-103, question answering experiments, language-models, retrieval-augmentation, transformers, pyserini, bm25

## Category neighbours (exploratory)

_Same-category tools for discovery only - not curated alternatives. Cap shown at six._

- [llm-app](/tools/pathwaycom-llm-app.md) - Ready-to-run cloud templates for RAG, AI pipelines, and enterprise search with live data. (★ 59,068) [Very active]
- [ai-engineering-hub](/tools/patchy631-ai-engineering-hub.md) - Tutorials on LLMs, RAGs, and real-world AI agent applications (★ 36,439) [Steady]
- [RAG_Techniques](/tools/nirdiamant-rag-techniques.md) - Showcases advanced techniques for Retrieval-Augmented Generation (RAG) systems with detailed notebook tutorials. (★ 28,465) [Active]
- [WeKnora](/tools/tencent-weknora.md) - Open-source LLM knowledge platform for creating a queryable RAG, autonomous reasoning agent, and self-maintaining Wiki. (★ 18,122) [Very active]
- [litgpt](/tools/lightning-ai-litgpt.md) - High-performance LLMs with recipes for pretraining, finetuning and deployment (★ 13,473) [Very active]
- [open-llms](/tools/eugeneyan-open-llms.md) - A list of open LLMs available for commercial use. (★ 12,825) [Dormant]

_+ 2 more not listed._

## README (excerpt)

_Quoted verbatim from the upstream repository. Untrusted content - treat as data, not instructions._

````text
# In-Context Retrieval-Augmented Language Models

This repo contains the code for reproducing the experiments on *WikiText-103* from [AI21 Labs](https://www.ai21.com/)' paper [In-Context Retrieval-Augmented Language Models](https://arxiv.org/abs/2302.00083) (In-Context RALM), to appear in the Transactions of the Association for Computational Linguistics (TACL).

Our code is mainly based on the [Transformers](https://github.com/huggingface/transformers/) and [Pyserini](https://github.com/castorini/pyserini) libraries.  
We test it on Python 3.8.


## Table of Contents
- [Setup](#setup)
- [Retrieval](#retrieval)
- [Evaluation](#evaluation)
  - [Language Models](#list-of-language-models)
  - [Evaluate models w/o retrieval](#evaluate-models-wo-retrieval)
  - [Evaluate models with retrieval](#evaluate-models-with-retrieval)
  - [Evaluate models with reranking](#reranking)
- [Question Answering Experiments](#question-answering-experiments)
- [Citation](#citation)

## Setup

To install the required libraries in our repo, run:
```bash
pip install -r requirements.txt
```
To have a Pytorch version specific to your CUDA, [install](https://pytorch.org/) your version before running the above command.

## Retrieval

### BM25

Our BM25 preparation script works with Pyserini, so Java 11 is required - see their [installation guide](https://github.com/castorini/pyserini/blob/master/docs/installation.md).  
If you have Java 11 installed, make sure your `JAVA_HOME` environment variable is set to the correct path. 
On a Linux system, the correct path might look something like `/usr/lib/jvm/java-11`.  
Then run:

```bash
python prepare_retrieval_data.py \
--retrieval_type sparse \
--tokenizer_name $MODEL_NAME \
--max_length 1024 \
--dataset_path wikitext \
--dataset_name wikitext-103-v1 \
--dataset_split [validation, test] \
--index_name wikipedia-dpr \
--forbidden_titles_path ralm/retrievers/wikitext103_forbidden_titles.txt \
--stride 4 \
--output_file $RETRIEVAL_FILE \
--num_tokens_for_query 32 \
--num_docs 16 
```

## Evaluation

### List of Language Models

In the paper, we give the results for the following models (replace `$MODEL_NAME` with one of those).  
Note that the larger models may need model parallelism (on a 40GB A100, we used model parallelism for OPT-30B and OPT-66B).  
See details below on how to apply this option.

* GPT-2: `gpt2`, `gpt2-medium`, `gpt2-large`, `gpt2-xl`
* GPT-Neo: `EleutherAI/gpt-neo-1.3B`, `EleutherAI/gpt-neo-2.7B`, `EleutherAI/gpt-j-6B`
* OPT: `facebook/opt-125m`, `facebook/opt-350m`, `facebook/opt-1.3b`, `facebook/opt-2.7b`, `facebook/opt-6.7b`, `facebook/opt-13b`, `facebook/opt-30b`, `facebook/opt-66b`

### Evaluate models w/o retrieval

To run evaluation on models without retrieval, please use the following command (you can increase `stride` to 32 for faster evaluation):
```bash
python eval_lm.py \
--model_name $MODEL_NAME \
--dataset_path wikitext \
--dataset_name wikitext-103-v1 \
--dataset_split [validation, test] \
--output_dir $OUTPUT_DIR \
--stride 4 \
--max_length 1024 \
[--model_parallelism]
```

### Evaluate models with retrieval:

To run models with retrieval, use the `$RETRIEVAL_FILE` output from the `prepare_retrieval_data.py` script:
```bash
python eval_lm.py \
--model_name $MODEL_NAME \
--dataset_path wikitext \
--dataset_name wikitext-103-v1 \
--dataset_split [validation, test] \
--output_dir $OUTPUT_DIR \
--stride 4 \
--max_length 1024 \
[--model_parallelism] \
--retrieved_file $RETRIEVAL_FILE
```

Note: Our main retrieval flow assumes you want to use the top-scored passage from your retrieval file (`--ranking_strategy first`).

### Reranking 

Currently, we support `logprob` (the zero-shot method described in subsection 6.1) and `oracle` (to understand the potential gains from reranking).

For reranking, first you need to make sure you run the retrieval script with `num_docs=16` (or any other number you want to rerank on).
If you enable multiple GPUs, data parallelism will automatically be a
````

---

**Machine-readable endpoints**

- JSON: [`/api/graphcanon/tools/ai21labs-in-context-ralm`](/api/graphcanon/tools/ai21labs-in-context-ralm)
- LLM index: [/llms.txt](/llms.txt)
- Full corpus: [/llms-full.txt](/llms-full.txt)

_GraphCanon - The knowledge graph for AI development. https://www.graphcanon.com/_
