From ee8a9fc5bc81fc889562a39e8e2357654efd5304 Mon Sep 17 00:00:00 2001 From: emil Date: Sat, 16 May 2026 13:18:31 +0300 Subject: [PATCH] Fix cross-process recall: MEMB trailer + raw eval/sample in chat() MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- experiments/README.md | 63 ++++++++++++------ experiments/diag_nemotron.py | 95 ++++++++++++++++++++++++++++ experiments/diag_nemotron2.py | 79 +++++++++++++++++++++++ experiments/diag_nemotron3.py | 91 ++++++++++++++++++++++++++ experiments/diag_session_nemotron.py | 32 ++++++++++ experiments/mood_batch_poc.py | 2 +- experiments/mood_stream_poc.py | 2 +- experiments/recall_poc.py | 2 +- python/memba/core.py | 29 ++++++++- python/memba/session.py | 63 ++++++++++++++---- 10 files changed, 424 insertions(+), 34 deletions(-) create mode 100644 experiments/diag_nemotron.py create mode 100644 experiments/diag_nemotron2.py create mode 100644 experiments/diag_nemotron3.py create mode 100644 experiments/diag_session_nemotron.py diff --git a/experiments/README.md b/experiments/README.md index 9a5289f..f382125 100644 --- a/experiments/README.md +++ b/experiments/README.md @@ -1,32 +1,57 @@ # experiments/ -Throwaway scripts used to probe capabilities of SSM models with memba. +Throwaway scripts used to probe SSM/hybrid model capabilities with memba. Not part of the library API — kept in the repo as reference and reproducible evidence for product decisions. -Each script is self-contained and prints what it finds; read the source for -the test's claim and run it yourself if you want to verify on different -models or hardware. +Each script is self-contained. Read the source for what it claims to test +and run it yourself if you want to verify on different models or hardware. ## Scripts | File | What it measures | |------|------------------| -| `recall_poc.py` | Can a memba state, built from N days of git activity, answer "what did I work on last month" in a fresh process? | -| `mood_poc.py` | Batch sentiment-trajectory test (single prompt with full chat log). Compares Falcon-Mamba vs a Transformer. | -| `mood_stream_poc.py` | The same trajectory but fed turn-by-turn through `Session.chat()`, then queried cross-process. | -| `mood_batch_poc.py` | Batch ingest in build process, save, then query in a fresh process. | -| `diag_saveload.py` | Minimal diagnostic: tell the model one fact, ask it back before save, after save, after cross-process load. | +| `recall_poc.py` | git log → state → cross-process "what did I work on last month" | +| `mood_poc.py` | Batch sentiment trajectory in one prompt — Falcon-Mamba vs Gemma | +| `mood_batch_poc.py` | Batch ingest, save, query in fresh process | +| `mood_stream_poc.py` | 15 separate observation turns, save, query in fresh process | +| `diag_saveload.py` | Minimal hamster recall test on Falcon-Mamba | +| `diag_nemotron.py` | Same hamster test on Nemotron 4B hybrid | +| `diag_nemotron2.py` | Cross-process recall via raw `generate()` (bypass `create_chat_completion`) | +| `diag_nemotron3.py` | Same with llama-cpp-python's native `save_state()`/`load_state()` | +| `diag_session_nemotron.py` | Full hamster recall through the rewritten `Session.chat()` | -## Headline finding (2026-05-16, Falcon-Mamba-7B-Instruct Q4_K_M) +## Findings log -- **Batch single-prompt analysis** (all input + question in one call): works - for both sentiment and recall. -- **Multi-turn fact recall** (ingest in turn 1, ask in turn 2): fails even - in the *same process*. The model does not preserve specific facts in its - hidden state across conversation turns. -- **Save/load roundtrip**: byte-identical, no information loss attributable - to memba's file format. The persistence layer works correctly; the - trained model just doesn't use the state for cross-turn recall. +### 2026-05-16 (initial, Falcon-Mamba-7B-Instruct Q4_K_M) -See the script outputs (or rerun) for the raw evidence. +- Batch single-prompt analysis works for both sentiment and recall. +- Multi-turn fact recall failed even in-process. *Misdiagnosed at first as + a model capability issue.* + +### 2026-05-16 (revised, after Nemotron-3-Nano-4B Q4_K_M test) + +- Real root cause: `Llama.create_chat_completion()` and `Llama.generate()` + retokenise the entire prompt each call and reset KV-cache when the new + tokens don't prefix-match `input_ids`. Loaded memba state was being wiped. +- Secondary issue: `llama_state_set_data()` restores the C-level cache but + llama-cpp-python's wrapper still reports `n_tokens=0`, so the next + `eval()` writes new tokens at offset 0 and overwrites the loaded state. +- Two fixes in memba: + 1. MEMB file now carries a Python-only trailer with `n_tokens` so the + wrapper position is restored after `load_state` (12 extra bytes, + backward-compatible — the C library ignores anything after the CRC). + 2. `Session.chat()` rewritten to use raw `tokenize → eval → sample` + instead of `Llama.__call__`, avoiding the prefix-matching reset. +- Results on Nemotron-3-Nano-4B (hybrid: 21 Mamba-2 + 4 attention layers) + through the new `Session.chat()`: + - `diag_session_nemotron.py` — cross-process hamster recall: ✅ + - `mood_batch_poc.py` — cross-process mood trajectory: ✅ + - `mood_stream_poc.py` — 15 streaming turns + cross-process: ✅ + - `recall_poc.py` — per-project digest from 30-day git log: ✅ accurate, + no hallucinated projects. + +Falcon-Mamba was rerun through the new `Session.chat()` but still fails on +the same tests — the trained model genuinely doesn't have the cross-turn +recall capability that the hybrid Nemotron does. The 4 attention layers +make the difference. diff --git a/experiments/diag_nemotron.py b/experiments/diag_nemotron.py new file mode 100644 index 0000000..67d69d1 --- /dev/null +++ b/experiments/diag_nemotron.py @@ -0,0 +1,95 @@ +""" +diag_nemotron.py — hamster recall test on Nemotron 3 Nano 4B (hybrid Mamba-Transformer). + +The Falcon-Mamba 0/4 hamster failure was the killshot for several product ideas. +This rerun tests whether the hybrid architecture (21 Mamba-2 layers + 4 attention) +fixes cross-turn recall. + +Three scenarios are measured: + 1. In-process multi-turn (tell fact, ask next turn) + 2. Same process: save then ask after save + 3. Cross process: build (ingest, save, exit), then query (load, ask) + +Uses create_chat_completion which applies the GGUF's own chat template. +""" +from __future__ import annotations +import sys, argparse, time +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) +from llama_cpp import Llama +from memba import core + +MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf" +STATE = "/tmp/diag_nemotron.memb" + +INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named " + "Bartholomew. He is 4 years old. Reply with just 'noted'.") +QUERY = "What is the name of my pet?" + + +def make_llama(): + return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False) + + +def chat_continued(m, messages): + """Send accumulated chat history, return assistant text + cleaned (strip reasoning).""" + out = m.create_chat_completion( + messages=messages, + max_tokens=200, + temperature=0.1, + ) + full = out["choices"][0]["message"]["content"].strip() + # Nemotron leaks reasoning — try to extract the final answer if present + short = full[-300:] if len(full) > 300 else full + return full, short + + +def build(): + print(f"[build] loading Nemotron 4B…") + t0 = time.time() + m = make_llama() + print(f"[build] loaded in {time.time()-t0:.1f}s") + + messages = [{"role": "user", "content": INGEST}] + full, _ = chat_continued(m, messages) + print(f"[build] ack (full):\n{full!r}\n") + messages.append({"role": "assistant", "content": full}) + + # Test 1: in-process recall WITHIN the same chat + messages.append({"role": "user", "content": QUERY}) + full, short = chat_continued(m, messages) + print(f"[build] in-process query BEFORE save:\n{full}\n") + messages.append({"role": "assistant", "content": full}) + + # Save state at this point + core.save_state(m, MODEL, STATE) + print(f"[build] saved state ({Path(STATE).stat().st_size:,} B)") + + # Test 2: in-process query AFTER save — should still work + messages.append({"role": "user", "content": QUERY}) + full, _ = chat_continued(m, messages) + print(f"[build] in-process query AFTER save:\n{full}\n") + + +def query(): + if not Path(STATE).exists(): + print("[query] no state — run build first"); return 1 + print(f"[query] loading model + state ({Path(STATE).stat().st_size:,} B)…") + t0 = time.time() + m = make_llama() + core.load_state(m, MODEL, STATE) + print(f"[query] loaded in {time.time()-t0:.1f}s") + + # Cross-process: send a fresh user turn with only the question + # The state should already encode the prior conversation + out = m.create_chat_completion( + messages=[{"role": "user", "content": QUERY}], + max_tokens=200, + temperature=0.1, + ) + print(f"[query] cross-process answer:\n{out['choices'][0]['message']['content']}") + + +if __name__ == "__main__": + cmd = sys.argv[1] if len(sys.argv) > 1 else "build" + {"build": build, "query": query}[cmd]() diff --git a/experiments/diag_nemotron2.py b/experiments/diag_nemotron2.py new file mode 100644 index 0000000..98374a8 --- /dev/null +++ b/experiments/diag_nemotron2.py @@ -0,0 +1,79 @@ +""" +diag_nemotron2.py — same test, but in cross-process query phase we use +raw __call__ with reset=False to NOT clear KV-cache before generation. + +Hypothesis: create_chat_completion() resets KV-cache on each call, which +defeats memba's loaded state. Bypassing that should let state work. +""" +import sys, time +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) +from llama_cpp import Llama +from memba import core + +MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf" +STATE = "/tmp/diag_nemotron.memb" # reuse state from previous test + +def make_llama(): + return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False) + + +def query_raw(): + """Cross-process query, bypassing create_chat_completion's reset.""" + if not Path(STATE).exists(): + print("[query] no state — run diag_nemotron.py build first") + return + print(f"[query] loading state ({Path(STATE).stat().st_size:,} B)…") + t0 = time.time() + m = make_llama() + print(f"[query] model loaded in {time.time()-t0:.1f}s") + + core.load_state(m, MODEL, STATE) + print(f"[query] state loaded in {time.time()-t0:.1f}s total") + + # Continue the conversation: just append a new user turn manually. + # Nemotron uses ChatML-like markers based on what we've seen: + # <|im_start|>user\n…<|im_end|>\n<|im_start|>assistant\n + # If wrong, model will still answer something interpretable. + continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n" + + # Tokenize continuation WITHOUT BOS — we're mid-sequence + tokens = m.tokenize(continuation.encode("utf-8"), add_bos=False, special=True) + print(f"[query] continuation tokenized to {len(tokens)} tokens") + + # Use generate() which feeds via eval (does NOT reset KV-cache) + print(f"[query] generating tokens (preserving loaded state)…") + out_tokens = [] + eos = m.token_eos() + for token in m.generate(tokens, top_k=1, temp=0.0): # greedy + if token == eos: + break + out_tokens.append(token) + if len(out_tokens) >= 200: + break + # Also stop at <|im_end|> + snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore") + if "<|im_end|>" in snippet: + break + + text = m.detokenize(out_tokens).decode("utf-8", errors="ignore") + print(f"\n[query] answer:\n{text}") + + +def query_reset_default(): + """Control: same continuation but with default reset=True (clears KV).""" + if not Path(STATE).exists(): + return + print(f"\n[control] same continuation, default reset=True (should fail):") + m = make_llama() + core.load_state(m, MODEL, STATE) + out = m( + "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n", + max_tokens=120, temperature=0.1, echo=False, # reset=True is default + ) + print(out["choices"][0]["text"]) + + +if __name__ == "__main__": + query_raw() + # query_reset_default() # optionally enable for direct comparison diff --git a/experiments/diag_nemotron3.py b/experiments/diag_nemotron3.py new file mode 100644 index 0000000..85d40cc --- /dev/null +++ b/experiments/diag_nemotron3.py @@ -0,0 +1,91 @@ +""" +diag_nemotron3.py — test hypothesis that llama-cpp-python's built-in +save_state/load_state preserves Python-side trackers (n_tokens, input_ids) +which memba's raw C-level save/load is missing. + +If built-in works → memba's MEMB format needs to be extended to include +those trackers. If built-in also fails → the issue is elsewhere. +""" +import sys, time, pickle +from pathlib import Path +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) +from llama_cpp import Llama + +MODEL = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf" +STATE = "/tmp/diag_nemotron_native.pickle" + +INGEST = ("I'm going to tell you a fact about my pet. My pet hamster is named " + "Bartholomew. He is 4 years old. Reply with just 'noted'.") +QUERY = "What is the name of my pet?" + + +def make_llama(): + return Llama(model_path=MODEL, n_ctx=4096, n_gpu_layers=-1, verbose=False) + + +def build(): + m = make_llama() + print(f"[build] initial n_tokens={m.n_tokens}") + + msgs = [{"role": "user", "content": INGEST}] + out = m.create_chat_completion(messages=msgs, max_tokens=80, temperature=0.1) + ack = out["choices"][0]["message"]["content"] + msgs.append({"role": "assistant", "content": ack}) + print(f"[build] after ingest: n_tokens={m.n_tokens}") + print(f"[build] ack snippet: {ack[-100:]!r}") + + # Use llama-cpp-python's NATIVE save_state — captures Python trackers too + state = m.save_state() + with open(STATE, "wb") as f: + pickle.dump(state, f) + print(f"[build] saved native state to {STATE} ({Path(STATE).stat().st_size:,} B)") + + +def query(): + if not Path(STATE).exists(): + print("[query] no state file"); return + m = make_llama() + print(f"[query] before load: n_tokens={m.n_tokens}") + + with open(STATE, "rb") as f: + state = pickle.load(f) + m.load_state(state) + print(f"[query] after load: n_tokens={m.n_tokens}") + + # Now ask via create_chat_completion. Pass ONLY the new question + # (history is already in KV-cache+n_tokens). + # The template will format this as if it's turn 1 — see what happens. + out = m.create_chat_completion( + messages=[{"role": "user", "content": QUERY}], + max_tokens=150, temperature=0.1, + ) + print(f"\n[query] answer via chat_completion (resets context):\n{out['choices'][0]['message']['content']}") + + # Alternative: try to continue manually, using raw eval + print(f"\n[query] re-loading state for raw continuation test") + with open(STATE, "rb") as f: + state = pickle.load(f) + m.load_state(state) + print(f"[query] re-loaded: n_tokens={m.n_tokens}") + + # Append a new user turn via raw tokens, then generate + continuation = "<|im_end|>\n<|im_start|>user\nWhat is the name of my pet?<|im_end|>\n<|im_start|>assistant\n" + tokens = m.tokenize(continuation.encode(), add_bos=False, special=True) + + out_tokens = [] + eos = m.token_eos() + for tok in m.generate(tokens, top_k=1, temp=0.0): + if tok == eos or len(out_tokens) >= 150: + break + out_tokens.append(tok) + snippet = m.detokenize(out_tokens).decode("utf-8", errors="ignore") + if "<|im_end|>" in snippet: + break + + text = m.detokenize(out_tokens).decode("utf-8", errors="ignore") + print(f"\n[query] raw continuation answer:\n{text}") + + +if __name__ == "__main__": + cmd = sys.argv[1] if len(sys.argv) > 1 else "build" + {"build": build, "query": query}[cmd]() diff --git a/experiments/diag_session_nemotron.py b/experiments/diag_session_nemotron.py new file mode 100644 index 0000000..ef9205b --- /dev/null +++ b/experiments/diag_session_nemotron.py @@ -0,0 +1,32 @@ +"""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]]() diff --git a/experiments/mood_batch_poc.py b/experiments/mood_batch_poc.py index b064186..9231510 100644 --- a/experiments/mood_batch_poc.py +++ b/experiments/mood_batch_poc.py @@ -11,7 +11,7 @@ from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) from memba import Session -MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf" +MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf" STATE_DIR = "/tmp/mood_batch_test" SESSION = "mood_batch" diff --git a/experiments/mood_stream_poc.py b/experiments/mood_stream_poc.py index 502807f..77a970e 100644 --- a/experiments/mood_stream_poc.py +++ b/experiments/mood_stream_poc.py @@ -14,7 +14,7 @@ from pathlib import Path sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) from memba import Session -MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf" +MAMBA = "/home/emil/Desktop/Coding/AI/Memba/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf" STATE_DIR = "/tmp/mood_stream_test" SESSION = "mood_stream" diff --git a/experiments/recall_poc.py b/experiments/recall_poc.py index 67da52a..d9dd95f 100644 --- a/experiments/recall_poc.py +++ b/experiments/recall_poc.py @@ -34,7 +34,7 @@ SCAN_ROOTS = [ ] DAYS_BACK = 30 STATE_DIR = Path.home() / ".recall/states" -SESSION_ID = "recall_poc" +SESSION_ID = "recall_poc_nemotron" DEFAULT_Q = ( "Summarise what I've been working on over the last month. " "Group by project. Mention the main themes per project. " diff --git a/python/memba/core.py b/python/memba/core.py index 5208065..1bf6671 100644 --- a/python/memba/core.py +++ b/python/memba/core.py @@ -25,6 +25,16 @@ MODEL_ID_LEN = 64 _HEADER_FMT = "<4sI64sIIQ" _HEADER_SIZE = struct.calcsize(_HEADER_FMT) # == 88 +# Optional Python trailer (appended after the V1 CRC32). The C library reads +# only up to data_size+CRC and ignores anything past it, so adding a trailer +# keeps the file C-compatible. The trailer carries llama-cpp-python's +# wrapper state (n_tokens) which is needed for `Llama.eval()` to append new +# tokens at the correct position instead of overwriting from offset 0. +_TRAILER_MAGIC = b"MTRL" +_TRAILER_FMT = "<4sII" # magic | version | n_tokens +_TRAILER_SIZE = struct.calcsize(_TRAILER_FMT) # == 12 +_TRAILER_VERSION = 1 + # ── Helpers ─────────────────────────────────────────────────────── @@ -140,9 +150,14 @@ def save_state(llama_model: "Llama", model_path: str, file_path: str) -> None: len(data), ) + # Python-only trailer with Llama wrapper position so eval() resumes + # at the right offset instead of overwriting from 0 after load_state. + n_tokens = getattr(llama_model, "n_tokens", 0) + trailer = struct.pack(_TRAILER_FMT, _TRAILER_MAGIC, _TRAILER_VERSION, n_tokens) + out = Path(file_path) out.parent.mkdir(parents=True, exist_ok=True) - out.write_bytes(header + data + struct.pack(" None: @@ -187,3 +202,15 @@ def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None: restored = _state_set_data(lib, ctx, buf, len(data)) if restored == 0: raise RuntimeError("llama_state_set_data returned 0 — state restore failed") + + # Optional trailer: restore llama-cpp-python wrapper position + trailer_off = offset + data_size + 4 + if len(raw) >= trailer_off + _TRAILER_SIZE: + magic, version, n_tokens = struct.unpack_from(_TRAILER_FMT, raw, trailer_off) + if magic == _TRAILER_MAGIC and version == _TRAILER_VERSION: + # Restore Llama wrapper state so eval() appends at the right + # offset instead of overwriting the loaded KV-cache from 0. + try: + llama_model.n_tokens = n_tokens + except AttributeError: + pass # wrapper version doesn't expose .n_tokens — eval may misbehave diff --git a/python/memba/session.py b/python/memba/session.py index cf37ccb..416f7e1 100644 --- a/python/memba/session.py +++ b/python/memba/session.py @@ -83,25 +83,66 @@ class Session: # ── Public methods ───────────────────────────────────────────── - def chat(self, prompt: str, max_tokens: int = 512) -> str: + def chat(self, prompt: str, max_tokens: int = 512, + temperature: float = 0.1, top_k: int = 1) -> str: """ Feed *prompt* to the model (wrapped in the active chat format) and return the generated text. - For SSM models the hidden state accumulates in llama_context across - calls — there is no explicit message history list, the recurrent - state IS the memory. Call save() at any checkpoint you want to - resume from later. + This uses raw tokenize+eval+sample rather than ``Llama.__call__`` or + ``create_chat_completion`` because those high-level helpers re-tokenise + the entire prompt and then reset the KV-cache when the prefix does not + match — which destroys any state loaded from a .memb file. Going + through ``llama.eval()`` directly *appends* new tokens to the live + state, which is exactly what memba needs. + + For SSM/hybrid models the hidden state accumulates in llama_context + across calls — there is no explicit message history list, the recurrent + state IS the memory. Call save() at any checkpoint you want to resume + from later. """ template, stops = CHAT_FORMATS[self._chat_format] wrapped = template.format(prompt=prompt) - result = self._llama( - wrapped, - max_tokens=max_tokens, - echo=False, - stop=stops, + + # Tokenise the wrapped turn. Add BOS only when the context is fresh + # (no prior tokens in either the live conversation or a loaded state). + is_fresh = self._llama.n_tokens == 0 + new_tokens = self._llama.tokenize( + wrapped.encode("utf-8"), + add_bos=is_fresh, + special=True, ) - return result["choices"][0]["text"].strip() + + # Append to the live state. eval() does NOT reset the KV-cache. + self._llama.eval(new_tokens) + + # Build the set of single-token stops + add EOS. + stop_token_ids = {self._llama.token_eos()} + for s in stops: + for t in self._llama.tokenize(s.encode("utf-8"), add_bos=False, special=True): + stop_token_ids.add(t) + + out_tokens: list[int] = [] + for _ in range(max_tokens): + tok = self._llama.sample(top_k=top_k, temp=temperature) + out_tokens.append(tok) + # Always eval the sampled token so it lives in the state too — + # that way the next chat() turn sees the assistant reply as + # part of the conversation. + self._llama.eval([tok]) + if tok in stop_token_ids: + break + # Multi-token stop-string check (some stops span several BPE pieces) + if stops: + snippet = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore") + if any(s in snippet for s in stops): + break + + text = self._llama.detokenize(out_tokens).decode("utf-8", errors="ignore") + for s in stops: + if s in text: + text = text.split(s)[0] + return text.strip() def save(self, session_id: Optional[str] = None) -> Path: """Persist the current state. Returns the path written."""