# Co-LMLM: Continuous-Query Limited Memory Language Models

📄 arXiv • 🌐 Project page • 🤗 Models & Data

This repository contains the code accompanying the paper **"Co-LMLM: Continuous-Query Limited Memory Language Models."** It covers the full pipeline: data annotation, pretraining of the three model variants (**Co-LMLM** and the **Standard LM** / **LMLM-Asker** baselines), retrieval-index construction, and the perplexity / generation / downstream evaluation suite. Co-LMLM is a retrieval-aware pretraining method: factual spans in the training corpus are annotated as `span`, and the model learns to emit a continuous query vector at each `` position that retrieves the span's content from an external index at inference time, rather than memorizing it in-weights. The models in the paper are **SmolLM2 135M / 360M** trained **from scratch**; variants trained on FineWeb-Edu are also supported by the same configs. In the code, the Co-LMLM trainer and its package are under `colmlm/`; the two baselines are the **Standard LM** (`lm_baseline/`) and **LMLM-Asker** (`lmlm_asker/`). All runnable scripts use [`pyrallis`](https://github.com/eladrich/pyrallis) for configuration: every script is invoked with `--config_path path/to/config.yaml`, and the YAML is parsed into a top-level dataclass defined next to the script. The **main config class** referenced under each section is the authoritative source for which fields a YAML can contain — start there when building your own configs. Example configs are in a `configs/` directory next to each component. ## Contents - [Env Setup](#env-setup) - [Quick Start](#quick-start) - [Codebase Structure](#codebase-structure) - [1. Annotation Pipeline](#1-annotation-pipeline) - [2. Pretraining](#2-pretraining) - [3. Index Building](#3-index-building) - [4. Evaluation & Generation](#4-evaluation--generation) - [Released Artifacts](#released-artifacts) - [Citation](#citation) ## Env Setup Dependencies are managed with [`uv`](https://github.com/astral-sh/uv) (Python 3.12). The full dependency set is pinned in [`pyproject.toml`](pyproject.toml) / [`uv.lock`](uv.lock). To create a virtual environment and install everything: ```bash uv sync ``` All scripts are then invoked via `uv run`, e.g. ```bash PYTHONPATH=src uv run python -m lmlm.colmlm.train --config_path path/to/config.yaml ``` `PYTHONPATH=src` is required for the `python -m` / direct-script invocations because the source package layout is under `src/` (`src/annotation`, `src/lmlm`, plus the top-level `consts.py` / `common.py`). The example commands below set it explicitly. ### Output paths Global output paths (annotations, training outputs, indices, and W&B logs) are defined in [`src/consts.py`](src/consts.py). They are all rooted under `PROJECT_DIR`, which reads the `LMLM_PROJECT_DIR` environment variable and defaults to `~/lmlm-project`: ```bash export LMLM_PROJECT_DIR=/path/to/your/storage # default: ~/lmlm-project ``` ### API keys Several stages call hosted LLM APIs and read their keys from environment variables: - `GEMINI_API_KEY` — Gemini seed annotation and the SimpleQA / FActScore LLM graders (default provider). - `OPENAI_API_KEY` — the OpenAI-backed alternative for the SimpleQA / FActScore graders. ## Quick Start ```bash # Install (see Env Setup) and download the model + Wikipedia index (~113 GB, resumable). uv sync hf download lil-lab/CoLMLM-360M-FW --local-dir ./CoLMLM-360M-FW hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-wiki-index ./co-lmlm-wiki-index ``` ```python import sys; sys.path.insert(0, "src") # make the repo's packages importable (or set PYTHONPATH=src) from lmlm.eval.hf_generate import load_retriever_generator gen = load_retriever_generator( model_path="./CoLMLM-360M-FW", index_path="./co-lmlm-wiki-index", db_path="./co-lmlm-wiki-index/entries.db", max_new_tokens=48, ) result = gen.generate("The chemical symbol for the element gold is") print(result.text) # The chemical symbol for the element gold is Au Au. It is a transition metal # transition metal that is a slightly reddish-yellow slightly reddish-yellow metal. It is a # very dense metal and is used in jewelry, coins, and other items. It print(result.num_retrievals) # 3 for e in result.retrieved_entries: print(f"{e.text_value!r} (score {e.score:.3f})") # ' Au' (score 0.909) # ' transition metal' (score 0.873) # ' slightly reddish-yellow' (score 0.844) ``` Every time the model emits ``, it uses that position's hidden state to query the index and splices the retrieved value in as `value`, which the model then copies into its running text — so the facts (`Au`, `transition metal`, `slightly reddish-yellow`, each pulled from the Wikipedia index) come from the store, not the weights. The returned `result` also carries `retrieved_entries` (each a `SearchResult` with `.text_value` and `.score`), `num_retrievals`, `failed_retrievals`, and timing fields — inspect them to see exactly what was retrieved and where. To generate over a whole file of prompts (batch, non-interactive), use the vLLM script instead — see [§4.2 Generation](#42-generation). For the larger FineWeb-Edu + Wikipedia index and all other options, see [Released Artifacts](#released-artifacts) and [`generation.md`](src/lmlm/eval/generation.md). ## Codebase Structure ``` src/ ├── consts.py # Global output paths (rooted at LMLM_PROJECT_DIR) ├── common.py ├── annotation/ # Annotation pipeline (data side) │ ├── annotate/ # Annotators + example configs (gemini/mlm/question_hybrid) │ │ └── large_scale_annotation/ # Slurm job-array annotation at scale │ ├── training/ # Trainers for the MLM and question generator │ │ ├── mlm/ # Fact-span detector (token classification) │ │ └── question_generator/ # Per-span question generator (causal LM) │ ├── data_selection/ # Subsetting / preparing source corpora │ └── prompts/ # Versioned Gemini prompts └── lmlm/ # Modeling, indexing, and evaluation ├── colmlm/ # Co-LMLM trainer ├── lmlm_asker/ # LMLM-Asker baseline trainer ├── lm_baseline/ # Standard LM baseline trainer ├── index/ # Index building (asker / retriever) │ └── large_scale_index/ # Sharded, Slurm-based index construction └── eval/ # Perplexity, dynamic replacement, generation, factuality, retrieval scripts/ └── eval/ # NLU eval, inference-efficiency, results collection ``` --- ## 1. Annotation Pipeline The pipeline produces a pretraining corpus where factual spans are annotated as `span` (`q` is the question, `a` is a paraphrased answer). There are three stages: (1) generate ground-truth annotations with Gemini on a seed set, (2) train smaller annotators on that seed, and (3) run the trained annotators over the full corpus. An overview of all annotator types is in [`src/annotation/OVERVIEW.md`](src/annotation/OVERVIEW.md). Source-corpus subsetting is handled by [`src/annotation/data_selection/data_selection.py`](src/annotation/data_selection/data_selection.py) (main config `DataSelectionConfig`; example [`configs/initial-5k-set.yaml`](src/annotation/data_selection/configs/initial-5k-set.yaml)). ### 1.1 Gemini seed annotation Gemini produces the ground-truth annotations for the seed set. The script is [`annotate_with_gemini.py`](src/annotation/annotate/annotate_with_gemini.py), and the versioned prompts (`system.txt`, `prompt.txt`, optional `continue_prompt.txt`) are under [`src/annotation/prompts/`](src/annotation/prompts/) (released version: `claude_lmlm_opt_v2/v30.1`). ```bash PYTHONPATH=src uv run src/annotation/annotate/annotate_with_gemini.py \ --config_path src/annotation/annotate/configs/gemini/dolmino-wiki.yaml ``` **Main config class:** `GeminiAnnotationConfig` in [`annotate_with_gemini.py`](src/annotation/annotate/annotate_with_gemini.py). It supports three input sources (HuggingFace dataset, prepared directory, local JSONL), a `GeminiConfig` selecting the model/prompt version, and a sharded on-disk cache so partial runs resume. ### 1.2 Training the annotators Three annotator components are trained from the Gemini-generated seed. Each is launched through a `run_training.sh ` wrapper around `accelerate launch`. **MLM fact-span annotator** — a ModernBERT token-classification head predicting the `` / `` span boundaries. ```bash bash src/annotation/training/mlm/run_training.sh \ src/annotation/training/mlm/configs/modernbert_large.yaml ``` Main config class: `MLMExperimentConfig` in [`src/annotation/training/mlm/config.py`](src/annotation/training/mlm/config.py). Supports LoRA or full fine-tuning, multi-GPU via `accelerate`, and seqeval / span-IoU metrics. Example config: [`modernbert_large.yaml`](src/annotation/training/mlm/configs/modernbert_large.yaml). **Question generator** — a causal LM trained to emit the question for a single fact span, given the surrounding annotated context. ```bash bash src/annotation/training/question_generator/run_training.sh \ src/annotation/training/question_generator/configs/default.yaml ``` Main config class: `QuestionGeneratorExperimentConfig` in [`src/annotation/training/question_generator/config.py`](src/annotation/training/question_generator/config.py). Backed by TRL's `SFTTrainer`, with optional config-controlled question-leakage filtering, document splitting, and question packing (efficient multi-question-per-document training). ### 1.3 Running the annotators For the paper we use the **question-hybrid** annotator: the MLM produces the fact spans, then the question generator fills in a question per span using a shared, prefilled KV cache for efficiency. ```bash CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \ src/annotation/annotate/annotate_with_question_hybrid.py \ --config_path src/annotation/annotate/configs/question_hybrid/example.yaml ``` Main config class: `QuestionHybridAnnotationConfig` in [`annotate_with_question_hybrid.py`](src/annotation/annotate/annotate_with_question_hybrid.py). Supports five input modes (HF dataset, prepared dir, an already-annotated corpus whose questions should be regenerated, local JSON, local JSONL). The MLM-only annotator is also available for span detection on its own: ```bash CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \ src/annotation/annotate/annotate_with_mlm.py \ --config_path src/annotation/annotate/configs/mlm/dolmino-wiki.yaml ``` Main config class: `MLMAnnotationConfig` in [`annotate_with_mlm.py`](src/annotation/annotate/annotate_with_mlm.py). ### 1.4 Annotation at scale For very large corpora there is a Slurm-based job-array runner under [`large_scale_annotation/`](src/annotation/annotate/large_scale_annotation/); the entry point is [`run_manager.py`](src/annotation/annotate/large_scale_annotation/run_manager.py) with top-level config `LargeScaleAnnotationConfig`. ```bash PYTHONPATH=src uv run -m annotation.annotate.large_scale_annotation.run_manager \ --config_path src/annotation/annotate/large_scale_annotation/configs/example_config.yaml ``` --- ## 2. Pretraining All three pretraining variants share the same streaming-dataset machinery (see [`src/lmlm/dataset_base.py`](src/lmlm/dataset_base.py) and the per-variant `dataset.py` files) and are launched the same way: `uv run python -m ` for single-GPU, or `accelerate launch -m ` for multi-GPU. Each loads an annotated corpus produced by the annotation pipeline above. The released models are SmolLM2 135M / 360M trained from scratch; FineWeb-Edu variants use the same configs with a different data source. ### 2.1 Co-LMLM The retrieval-aware pretraining method introduced in the paper. The trainer masks out fact-span content from the next-token-prediction loss and adds an InfoNCE fact–question contrastive objective so that `` hidden states become useful continuous retrieval queries. ```bash PYTHONPATH=src uv run accelerate launch -m lmlm.colmlm.train \ --config_path src/lmlm/colmlm/configs/example_config.yaml ``` Main config class: `LMLMConfig` in [`src/lmlm/colmlm/config.py`](src/lmlm/colmlm/config.py). Sub-configs to look at: `ModelConfig`, `ContrastiveLossConfig`, `DataConfig`, `TrainingConfig`, `OptimizerConfig`. An in-loop retrieval-eval example (retrieval metrics + full-eval perplexity computed during training) is provided in [`example_inloop_retrieval_eval.yaml`](src/lmlm/colmlm/configs/example_inloop_retrieval_eval.yaml). ### 2.2 Standard LM (baseline) A plain causal-LM baseline trained on the *unannotated* text (fact tags are stripped). Loss is broken out per token role (overall / fact-span / rest) so the baseline's perplexity on fact spans can be compared directly to the LMLM variants. ```bash PYTHONPATH=src uv run accelerate launch -m lmlm.lm_baseline.train \ --config_path src/lmlm/lm_baseline/configs/example_config.yaml ``` Main config class: `LMBaselineConfig` in [`src/lmlm/lm_baseline/config.py`](src/lmlm/lm_baseline/config.py). ### 2.3 LMLM-Asker (baseline) The "ask, then retrieve" baseline: the model is trained to emit `question` *before* each fact span (the question sits between `` and ``), so at inference the natural-language question is used as a query against an external sentence-transformer-encoded index. The fact-span content itself is excluded from the loss. ```bash PYTHONPATH=src uv run accelerate launch -m lmlm.lmlm_asker.train \ --config_path src/lmlm/lmlm_asker/configs/example_config.yaml ``` Main config class: `LMLMAskerConfig` in [`src/lmlm/lmlm_asker/config.py`](src/lmlm/lmlm_asker/config.py). --- ## 3. Index Building Inference for both LMLM-Asker and Co-LMLM requires a retrieval index over the *background corpus* (an annotated corpus, typically separate from the evaluation target). The single entry point is [`src/lmlm/index/build_index.py`](src/lmlm/index/build_index.py); exactly one of two sub-configs determines what gets built: - **`asker`** — encodes the questions (and stores their fact-span answers) with a `sentence-transformers` model into an `AskerIndex`. Used by LMLM-Asker. - **`retriever`** — runs a forward pass of the trained Co-LMLM model over the annotated background corpus and indexes the hidden state at each `` token (`RetrieverIndex`). This is the variant used for Co-LMLM in the paper. ```bash CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run -m lmlm.index.build_index \ --config_path path/to/your_index_config.yaml ``` Main config class: `IndexBuildConfig` in [`build_index.py`](src/lmlm/index/build_index.py) (with `IndexBuildSettings`, `AskerBuildConfig`, `RetrieverBuildConfig`). For corpora too large to embed in a single process, a sharded Slurm-based pipeline is in [`large_scale_index/`](src/lmlm/index/large_scale_index/) (see its [`README.md`](src/lmlm/index/large_scale_index/README.md)): `embedding_workers/` manages distributed embedding extraction, then `build_index.py` / `merge_sharded_index.py` build the FAISS index from the shards. The entry point is [`embedding_workers/run_manager.py`](src/lmlm/index/large_scale_index/embedding_workers/run_manager.py) with top-level config `LargeScaleIndexConfig`. --- ## 4. Evaluation & Generation The full evaluation flow is documented in [`src/lmlm/eval/eval_flow.md`](src/lmlm/eval/eval_flow.md) and the generation scripts (vLLM and HuggingFace backends) in [`src/lmlm/eval/generation.md`](src/lmlm/eval/generation.md). A summary follows. ### 4.1 Perplexity & dynamic replacement For LMLM-Asker and Co-LMLM the recommended path is the unified [`full_eval.py`](src/lmlm/eval/full_eval.py), which chains the four steps below into one script with auto-derived intermediate paths: 1. annotate the *target* and *background* datasets (using the annotators above); 2. compute static / normalized perplexity with [`compute_perplexity.py`](src/lmlm/eval/compute_perplexity.py); 3. build the background index (`lmlm.index.build_index`); 4. run dynamic replacement — [`vllm_asker_dynamic_replace.py`](src/lmlm/eval/vllm_asker_dynamic_replace.py) for LMLM-Asker, [`batched_retriever_dynamic_replace.py`](src/lmlm/eval/batched_retriever_dynamic_replace.py) for Co-LMLM — and recompute perplexity to obtain the dynamic-PPL number reported in the paper. ```bash CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/full_eval.py \ --config_path path/to/your_full_eval_config.yaml ``` Main config class: `FullEvalConfig` in [`full_eval.py`](src/lmlm/eval/full_eval.py). For the Standard LM baseline only step 2 is needed. Each step can also be run on its own — see [`eval_flow.md`](src/lmlm/eval/eval_flow.md) for per-step invocations. ### 4.2 Generation For **batch generation** over a file of prompts, use the **vLLM** scripts — one per model family, and the fastest for throughput. (For interactive, single-prompt use, prefer the `load_retriever_generator` helper from the [Quick Start](#quick-start); once you are running non-interactively vLLM is the better choice.) All scripts read JSONL prompts and write JSONL outputs; templates are in [`configs/generate/`](src/lmlm/eval/configs/generate/): | Model | vLLM script | Config template | |-------|-------------|-----------------| | Standard LM | [`vllm_generate.py`](src/lmlm/eval/vllm_generate.py) | [`lm_template.yaml`](src/lmlm/eval/configs/generate/lm_template.yaml) | | LMLM-Asker | [`vllm_asker_generate.py`](src/lmlm/eval/vllm_asker_generate.py) | [`asker_template.yaml`](src/lmlm/eval/configs/generate/asker_template.yaml) | | Co-LMLM | [`vllm_retriever_generate.py`](src/lmlm/eval/vllm_retriever_generate.py) | [`retriever_template.yaml`](src/lmlm/eval/configs/generate/retriever_template.yaml) | For LMLM-Asker, when `` is emitted the question is queried against an `AskerIndex` and the retrieved answer is injected (immediately followed by ``) before generation continues. For Co-LMLM, when `` is produced a separate HF model extracts the hidden state at that position; that vector queries a `RetrieverIndex` (retrieval runs concurrently with vLLM via a thread pool). ```bash # Fill in model_path / index_path / index.db_path / prompts_path / output_path in the template, # then run the Co-LMLM (retriever) generator: CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/vllm_retriever_generate.py \ --config_path src/lmlm/eval/configs/generate/retriever_template.yaml ``` The same retrieval loop is also exposed through the HuggingFace `transformers` backend — interactively via `load_retriever_generator(...)` ([Quick Start](#quick-start)) and as a config-driven script [`hf_generate.py`](src/lmlm/eval/hf_generate.py) (`--model_type asker|retriever`). See [`generation.md`](src/lmlm/eval/generation.md) for full invocation examples of both backends (including one- vs two-GPU setups for the retriever script). An eval-only baseline for the original **structured-LMLM** system (JSON knowledge base / FAISS lookup via explicit `<|db_entity|>` markup) is provided in [`vllm_lmlm_baseline_generate.py`](src/lmlm/eval/vllm_lmlm_baseline_generate.py) for comparison; its outputs feed the same downstream scorers below. (Training of structured-LMLM is out of scope for this repository.) ### 4.3 Downstream factuality Prompt-prep and scoring scripts sit next to the generation scripts. Prep writes JSONL prompts, a generation script produces continuations, then a scorer grades them. - **PopQA / SimpleQA** — [`prepare_popqa_prompts.py`](src/lmlm/eval/prepare_popqa_prompts.py) (`PrepareOpenQAPromptsConfig`, `--dataset_name popqa|simpleqa`) prepares both. PopQA is graded by [`score_popqa.py`](src/lmlm/eval/score_popqa.py) (`ScorePopQAConfig`, substring exact-match); SimpleQA is graded by the canonical LLM grader [`score_simpleqa.py`](src/lmlm/eval/score_simpleqa.py) (`ScoreSimpleQAConfig`, requires `GEMINI_API_KEY` or `OPENAI_API_KEY`). - **T-REx** — [`prepare_trex_prompts.py`](src/lmlm/eval/prepare_trex_prompts.py) (`PrepareTrexPromptsConfig`) and [`score_trex.py`](src/lmlm/eval/score_trex.py) (`ScoreTrexConfig`). - **FActScore** — [`prepare_factscore_prompts.py`](src/lmlm/eval/prepare_factscore_prompts.py) (`PrepareFactScorePromptsConfig`) prepares prompts; the atomic-fact scorer is in [`factscore/factscorer.py`](src/lmlm/eval/factscore/factscorer.py) (`FactScorerConfig`, requires an LLM API key). ```bash PYTHONPATH=src uv run src/lmlm/eval/prepare_popqa_prompts.py \ --dataset_name simpleqa \ --output_path output/eval/simpleqa/simpleqa_prompts.jsonl ``` ### 4.4 NLU evaluation Multiple-choice / language-understanding benchmarks via lighteval (with optional masking of LMLM factual special tokens). Driver: [`scripts/eval/eval_nlu_task.sh`](scripts/eval/eval_nlu_task.sh), which invokes [`scripts/eval/eval_nlu_masked.py`](scripts/eval/eval_nlu_masked.py). Task definitions are in [`src/lmlm/eval/lighteval_tasks.py`](src/lmlm/eval/lighteval_tasks.py). ```bash CUDA_VISIBLE_DEVICES=0 CKPT=/path/to/checkpoint SETTING=smollm2-setting \ bash scripts/eval/eval_nlu_task.sh ``` ### 4.5 RAG comparison baseline A retrieval-augmented-generation baseline that prepends retrieved passages to the prompt: [`prepare_rag_prompts.py`](src/lmlm/eval/prepare_rag_prompts.py) (`PrepareRagPromptsConfig`) builds the prompts and [`retrieve_rag_passages.py`](src/lmlm/eval/retrieve_rag_passages.py) (`RetrieveRagPassagesConfig`) attaches the retrieved passages; the resulting prompts feed the Standard LM generation script. ### 4.6 Inference efficiency Throughput / overhead measurements (decode rate, dynamic-PPL overhead, encoder rate) are under [`scripts/eval/inference_efficiency/`](scripts/eval/inference_efficiency/); see its [`README.md`](scripts/eval/inference_efficiency/README.md) for the `run_timed.sh` / `run_dynppl.sh` drivers and the analysis scripts. ### 4.7 Retrieval-precision eval Measures how often the correct entry is retrieved for a given query: [`eval_retrieval_precision.py`](src/lmlm/eval/eval_retrieval_precision.py) (`EvalRetrievalPrecisionConfig`). ```bash CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run src/lmlm/eval/eval_retrieval_precision.py \ --config_path src/lmlm/eval/configs/eval/retrieval_precision.yaml ``` --- ## Released Artifacts Released under the [**Co-LMLM** collection](https://huggingface.co/collections/lil-lab/co-lmlm-6a4e8216d55eae83af348f57) on the Hugging Face Hub (`lil-lab`). ### Co-LMLM-360M (FineWeb-Edu) - **Model** — [`lil-lab/CoLMLM-360M-FW`](https://huggingface.co/lil-lab/CoLMLM-360M-FW): the 360M SmolLM2-based Co-LMLM retriever trained on FineWeb-Edu (HF model repo: weights + tokenizer). Loaded with the standard HF loader; the fact special tokens (`` etc.) are in the tokenizer. Two retrieval indices are released for this model, both as Hugging Face [Storage Buckets](https://huggingface.co/docs/hub/storage-buckets) and both context-keyed on the model's `` hidden states. Each bucket holds the FAISS index (`faiss.index`), the faiss-id → entry-id mapping, `manifest.json` + `index_config.json`, and the fact-span value store (`entries.db`): | Index | Bucket | Background corpus | Entries | FAISS factory | Size | |-------|--------|-------------------|--------:|---------------|-----:| | **Wikipedia** | `co-lmlm-360m-fw-wiki-index` | full Wikipedia | 236M | `OPQ240,IVF65536,PQ240` | ~113 GB | | **FineWeb-Edu + Wikipedia** | `co-lmlm-360m-fw-fineweb-wiki-index` | FineWeb-Edu (100BT) + full Wikipedia | 2.2B | `OPQ96,IVF524288_HNSW32,PQ96` (~2.5× more compressed) | ~1.07 TB | Download the model and whichever index you need: ```bash # Model (weights + tokenizer) hf download lil-lab/CoLMLM-360M-FW --local-dir ./CoLMLM-360M-FW # Option A — Wikipedia index (~113 GB; resumable) hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-wiki-index ./co-lmlm-wiki-index # Option B — FineWeb-Edu + Wikipedia index (~1.07 TB; resumable) hf buckets sync hf://buckets/lil-lab/co-lmlm-360m-fw-fineweb-wiki-index ./co-lmlm-fineweb-wiki-index ``` Then point a Co-LMLM generation/eval config (or the [Quick Start](#quick-start) loader) at the model plus the chosen index: - **Wikipedia index** — `model_path: ./CoLMLM-360M-FW`, `index_path: ./co-lmlm-wiki-index`, `index.db_path: ./co-lmlm-wiki-index/entries.db`. - **FineWeb-Edu + Wikipedia index** — `model_path: ./CoLMLM-360M-FW`, `index_path: ./co-lmlm-fineweb-wiki-index`, `index.db_path: ./co-lmlm-fineweb-wiki-index/fineweb_with_fullwiki_entries.db`. This index ships its faiss-id → entry-id mapping as a SQLite database (`faiss_id_to_entry_id.db`) rather than a `.txt` file, so also set `index.use_sqlite_id_mapping: true`. The FineWeb-Edu + Wikipedia FAISS file alone is ~228 GB. If it does not fit in RAM, **memory-map** it instead of loading it fully by exporting `LMLM_FAISS_MMAP=1` — pages are then paged in on demand from local disk (lower RAM, slightly slower search): ```bash LMLM_FAISS_MMAP=1 CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src uv run \ src/lmlm/eval/hf_generate.py --model_type retriever \ --model_path ./CoLMLM-360M-FW \ --index_path ./co-lmlm-fineweb-wiki-index \ --index.db_path ./co-lmlm-fineweb-wiki-index/fineweb_with_fullwiki_entries.db \ --index.use_sqlite_id_mapping true \ --prompts_path src/lmlm/eval/configs/generate/prompts_example.jsonl --output_path out.jsonl ``` See [`retriever_template.yaml`](src/lmlm/eval/configs/generate/retriever_template.yaml) and [`generation.md`](src/lmlm/eval/generation.md). ### Other checkpoints & data The remaining checkpoints (135M / 360M Co-LMLM, Standard LM, and LMLM-Asker variants) and the annotated pretraining corpora will be added to the collection as they are released. --- ## Citation If you use this code or the released artifacts, please cite: ```bibtex @misc{feldman2026colmlmcontinuousquerylimitedmemory, title={Co-LMLM: Continuous-Query Limited Memory Language Models}, author={Yair Feldman and Linxi Zhao and Nathan Godey and Dongyoung Go and Yilun Hua and Kilian Q. Weinberger and Jennifer J. Sun and Yoav Artzi}, year={2026}, eprint={2607.07707}, archivePrefix={arXiv}, primaryClass={cs.CL}, url={https://arxiv.org/abs/2607.07707}, } ```