Files
Memba/experiments/mood_stream_poc.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

92 lines
3.5 KiB
Python

"""
mood_stream_poc.py — streaming sentiment, then save/load across processes.
This is the REAL product test:
1. Open memba session
2. Feed 15 chat messages ONE AT A TIME as "observations"
3. Save state, exit process
4. In a fresh process: load state, query mood
"""
from __future__ import annotations
import sys, argparse, time
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/NVIDIA-Nemotron3-Nano-4B-Q4_K_M.gguf"
STATE_DIR = "/tmp/mood_stream_test"
SESSION = "mood_stream"
CHAT_LOG = [
"Morning team! Coffee in hand, ready to tackle the auth refactor today.",
"Just pushed PR #234 fixing the token validation bug. Should be a quick merge.",
"Code review comments came in fast, all good catches. Iterating now.",
"Basic flow working locally, tests passing. Feeling good about this.",
"Heading to lunch, hopefully wrap this up by EOD.",
"Back. CI is failing on something unrelated, looking into it.",
"OK the 'unrelated' thing is actually related. Auth tests use a stale fixture.",
"Why does the fixture rebuild take 12 minutes. Every. Single. Time.",
"Cancelled the run twice now. Going to bypass and run tests locally.",
"Local passes, CI fails. Classic.",
"Two hours gone on this fixture issue. Not even what I was supposed to be doing.",
"Now there's a merge conflict with main because someone restructured migrations.",
"Whoever shipped those migrations on a Friday afternoon, I will find you.",
"Closing the laptop. Will fight this tomorrow.",
"Actually no. One more try before I sleep.",
]
QUESTIONS = [
"Briefly: what is this person's current emotional state? One sentence.",
"Has their mood changed during this monitoring session? One sentence.",
"Roughly when did they start having a hard time?",
]
def build():
state_path = Path(STATE_DIR) / f"{SESSION}.memb"
if state_path.exists():
state_path.unlink()
print(f"[build] cleared previous state")
s = Session(
model_path=MAMBA, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml",
)
print(f"[build] init state {s.state_size:,} B")
# Stream-feed: each message wrapped as if WE'RE TELLING the model
# "here's a new message you're observing"
for i, msg in enumerate(CHAT_LOG, 1):
observation = f"You are silently observing one person's chat messages. New message just arrived:\n[msg {i}] {msg}\nReply with just 'noted'."
ack = s.chat(observation, max_tokens=4)
print(f"[build] msg {i:>2}: {msg[:50]:<50} → ack={ack!r}")
print(f"[build] state after streaming: {s.state_size:,} B")
s.save()
print(f"[build] saved → {state_path}")
def query():
state_path = Path(STATE_DIR) / f"{SESSION}.memb"
if not state_path.exists():
print("[query] no state — run build first"); return 1
print(f"[query] loading state {state_path.stat().st_size:,} B")
s = Session(
model_path=MAMBA, session_id=SESSION, state_dir=STATE_DIR,
n_gpu_layers=-1, n_ctx=4096, chat_format="chatml",
)
for i, q in enumerate(QUESTIONS, 1):
print(f"\n[Q{i}] {q}")
print(f"[A{i}] {s.chat(q, max_tokens=120)}")
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("cmd", choices=["build", "query"])
args = p.parse_args()
if args.cmd == "build":
build()
else:
query()