Chimera (RLM x SSM)

Minimal PyTorch scaffold for an RLM-style orchestrator with an SSM chunk encoder and merge reducers.

Layout

• chimera/rlm/: prompt store, tool API, controller planner, Chimera LM wrapper
• chimera/ssm/: SSM-inspired encoder block and LM loss
• chimera/merge/: chunk-level reducers
• chimera/tasks/: PG-19-style dataset wrappers and baselines
• chimera/eval/: core metrics
• scripts/: train/eval/report entry points
• configs/: YAML starter configs

Quickstart

1. Install deps:

pip install -r requirements.txt

2. Prepare data:

• Place text files at:
  • data/pg19/train/*.txt
  • data/pg19/valid/*.txt (optional)

3. Train SSM baseline:

python scripts/train_pg19_ssm.py model=ssm data_dir=data/pg19 out=artifacts/train_ssm.json

4. Train Transformer baseline:

python scripts/train_pg19_ssm.py model=transformer data_dir=data/pg19 out=artifacts/train_transformer.json

5. Run long retrieval probe:

python scripts/eval_long_retrieval.py out=artifacts/retrieval_metrics.json

6. Build markdown report:

python scripts/report_tradeoffs.py out=artifacts/report.md

Notes

• Current SSM block is a lightweight SSM-inspired implementation for fast iteration.
• Swap in official Mamba/S4 kernels in chimera/ssm/encoder.py for production experiments.
• scripts/eval_long_retrieval.py uses synthetic needle probes for orchestration sanity checks.

Pull All Long-Context Benchmarks

To pull all benchmark repos (PG-19, BABILong, RULER, InfiniteBench, LongBench-v2) and bootstrap local data subsets:

bash scripts/pull_all_benchmarks.sh

To inspect registered benchmark paths:

python3 scripts/list_benchmarks.py

To customize data volume:

python3 scripts/prepare_all_benchmarks_data.py \
  --out-dir data/benchmarks \
  --pg19-train-max 1000 \
  --babilong-config 32k \
  --babilong-task qa1 \
  --infinitebench-max-lines-per-file 1000 \
  --longbench-max 500

Non-Blocking Training

Launch any long job in the background with logs + PID files:

bash scripts/launch_detached.sh ssm_train \
  python scripts/train_pg19_ssm.py model=ssm data_dir=data/benchmarks/pg19 steps=10000 out=artifacts/train_ssm.json

Check progress:

tail -f artifacts/logs/ssm_train.log

Privacy Audit Before Push

Run this before every public push:

bash scripts/privacy_audit.sh

If it returns non-zero, review and remove matched secrets/paths before pushing.

Lightning + Hydra Training

Install dependencies:

pip install -r requirements.txt

Run SSM with Hydra defaults:

python scripts/train_lightning.py

Override config at launch:

python scripts/train_lightning.py model=transformer trainer=gpu trainer.max_steps=2000 data.batch_size=2

Non-blocking launch:

bash scripts/launch_detached.sh lt_ssm \
  python scripts/train_lightning.py model=ssm trainer=gpu trainer.max_steps=10000