HiR

Minimal reference implementation of HiR, a training-free decoding method for calibrating how retrieved hierarchical memory influences long-term-memory question answering.

This release intentionally contains only the paper-facing LoCoMo pipeline:

• lexical retrieval over LoCoMo's summaries, observations, and dialogue turns;
• construction of coarse, mid, fine, and full memory views;
• greedy Full and adaptive HiR decoding;
• the official LoCoMo category-aware token-F1 evaluator.

It does not contain datasets, checkpoints, generated predictions, experiment results, server scripts, or exploratory variants.

Method

For every decoding step, the same frozen language model processes a full hierarchy view and a coarse-only reference view. HiR scores token v as

score(v) = log p(v | full) - alpha_t * log p(v | coarse).

The coefficient is adapted from the entropy gap between the two next-token distributions. The default alpha_0=0.65 reproduces the paper configuration.

Installation

Python 3.10 or newer is required.

python -m venv .venv
source .venv/bin/activate
pip install -e .

For development:

pip install -e '.[dev]'
pytest

Data

Download locomo10.json from the official LoCoMo repository. The data is not redistributed here; its original license and terms continue to apply.

Run

Both the dataset and output locations are explicit command-line arguments. No machine-specific paths are embedded in the code.

hir-locomo run \
  --data /path/to/locomo10.json \
  --model Qwen/Qwen3-8B \
  --output /path/to/predictions.jsonl \
  --methods full hir

Use --limit 10 for a smoke test. The main LoCoMo defaults are five retrieved sessions, twelve dialogue turns, five observations, a root timeline of eight sessions, 32 generated tokens, greedy decoding, and alpha_0=0.65.

Evaluate an existing prediction file with:

hir-locomo evaluate \
  --predictions /path/to/predictions.jsonl \
  --output /path/to/summary.json

Run hir-locomo run --help for model-loading and memory-construction options.

Output format

The runner writes one JSON object per example and method:

{"ex_id":"locomo:sample-0:0","category":1,"question":"...","answer":"...","method":"hir","prediction":"..."}

The evaluator reports overall and per-category official LoCoMo token-F1.

Reproducibility notes

• Generation is greedy (do_sample=False semantics).
• If a tokenizer provides a chat template, it is used by default.
• Qwen-style thinking is disabled when the tokenizer supports that option.
• The released hierarchy uses global leaf retrieval, matching the main LoCoMo setting in the paper.
• HiR requires token-level model outputs and therefore does not support text-only black-box generation APIs.