emilandClaude Opus 4.7
ee8a9fc5bc
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>
96 lines
3.4 KiB
Python
96 lines
3.4 KiB
Python
"""
|
|
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]()
|