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>
This commit is contained in:
+44
-19
@@ -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.
|
||||
|
||||
@@ -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]()
|
||||
@@ -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
|
||||
@@ -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]()
|
||||
@@ -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]]()
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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. "
|
||||
|
||||
+28
-1
@@ -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("<I", crc))
|
||||
out.write_bytes(header + data + struct.pack("<I", crc) + trailer)
|
||||
|
||||
|
||||
def load_state(llama_model: "Llama", model_path: str, file_path: str) -> 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
|
||||
|
||||
+52
-11
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user