Files
Memba/experiments/diag_session_nemotron.py
emilandClaude Opus 4.7 ee8a9fc5bc Fix cross-process recall: MEMB trailer + raw eval/sample in chat()
Two bugs were blocking memba's main promise (load .memb in a fresh
process → model continues with full recalled context):

1. llama_state_set_data() restores the C-level KV-cache + SSM hidden
   state, but llama-cpp-python's Python wrapper still reports
   n_tokens=0. The next eval() then decodes new tokens at offset 0
   and overwrites the loaded state.

   Fix: extend MEMB format with an optional 12-byte trailer appended
   after the CRC32. It carries the wrapper's n_tokens. The C library
   reads up to CRC and ignores anything past it, so files stay
   backward-compatible with libmemba; only the Python loader uses it.

2. Llama.__call__ / create_chat_completion / generate all re-tokenise
   the prompt on every call and clear the KV-cache when the new
   tokens don't prefix-match input_ids. That destroys any state we
   just loaded.

   Fix: rewrite Session.chat() to use raw tokenize → eval → sample.
   eval() appends tokens to the live state without resetting, and we
   handle stop-token detection ourselves.

Verified end-to-end on Nemotron-3-Nano-4B (hybrid 21x Mamba-2 + 4x
attention) — see experiments/README.md for the full findings log.
diag_session_nemotron.py, mood_batch_poc.py, mood_stream_poc.py and
recall_poc.py now all pass their cross-process tests; Falcon-Mamba
still fails because the trained model itself can't do cross-turn
recall — that was the original misdiagnosis.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
2026-05-16 13:18:31 +03:00

33 lines
1.3 KiB
Python

"""Same hamster test, but through the new Session.chat() (which uses eval+sample)."""
import sys, argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python"))
from memba import Session
MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/diag_session_nemotron"
SESSION = "hamster"
def build():
s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
# Wipe any previous state file so it's a real fresh run
p = Path(STATE_DIR) / f"{SESSION}.memb"
if p.exists(): p.unlink()
print(f"[build] ack: {s.chat('My pet hamster is named Bartholomew. He is 4 years old. Reply noted.', max_tokens=40)!r}")
print(f"[build] in-process query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
s.save()
print(f"[build] saved")
def query():
s = Session(model_path=MODEL, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml")
print(f"[query] loaded state, n_tokens={s._llama.n_tokens}")
print(f"[query] CROSS-PROCESS query: {s.chat('What is the name of my pet?', max_tokens=80)!r}")
if __name__ == "__main__":
{"build": build, "query": query}[sys.argv[1]]()