Add memba MVP: C++ core, Python SDK, CLI, examples, experiments

C++ core (libmemba.so):
- include/memba/state.h — C API (state_new/free/save/load/get_size)
- src/state.cpp — MEMB file format: magic, version, SHA-256 model_id,
  CRC-32, opaque llama_state_*_data() blob
- src/cli.cpp — minimal demo binary with greedy sampler
- CMakeLists.txt + build.sh with llama.cpp submodule, CUDA auto-detect

Python SDK (memba):
- core.py — file I/O via llama-cpp-python's exposed C functions,
  unwraps _LlamaContext to access raw context pointer (≥0.3.x)
- session.py — high-level Session with auto-save/load, ChatML wrapper
  for instruct models, raw mode for base models
- cli.py — typer-based: chat (REPL), run (one-shot), list, rm, info

Examples:
- 01_basic_save_load.py, 02_chat_session.py

Experiments (throwaway POCs documenting product-direction findings):
- recall_poc.py — git log → state → cross-process query
- mood_poc.py — batch sentiment trajectory, Mamba vs Transformer
- mood_stream_poc.py, mood_batch_poc.py — variants
- diag_saveload.py — minimal save/load isolation test
- README.md documents the headline finding: save/load is byte-identical,
  but Falcon-Mamba-7B-Instruct does not retain facts across conversation
  turns even in-process — limits viable products to single-prompt analysis
  and persona priming.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This commit is contained in:
emil
2026-05-16 12:48:37 +03:00
co-authored by Claude Opus 4.7
parent 3bb2d3dc71
commit 75c9ee4576
22 changed files with 2357 additions and 0 deletions
+143
View File
@@ -0,0 +1,143 @@
"""High-level Session API — wraps llama-cpp-python Llama + memba state I/O."""
from __future__ import annotations
from pathlib import Path
from typing import Optional
from . import core
# ── Chat formats ───────────────────────────────────────────────────
# Each entry maps a format name to (prompt_template, stop_strings).
# The template uses {prompt} as the placeholder for the user message.
CHAT_FORMATS: dict[str, tuple[str, list[str]]] = {
"chatml": (
"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n",
["<|im_end|>", "<|im_start|>"],
),
"raw": (
"{prompt}",
[],
),
}
class Session:
"""
A persistent SSM chat session backed by a memba state file.
The session auto-loads an existing state on construction (if one exists
for *session_id*) and accumulates context across calls to chat().
Call save() to persist the current state.
Parameters
----------
model_path: Path to the GGUF model file.
session_id: Logical name for this session; determines the state filename.
state_dir: Directory where .memb files are stored (created if absent).
n_gpu_layers: GPU layers to offload (0 = CPU-only, -1 = all layers).
n_ctx: Context window size in tokens.
verbose: Forward llama.cpp log output to stderr.
"""
def __init__(
self,
model_path: str,
session_id: str = "default",
state_dir: str = "~/.memba/states",
n_gpu_layers: int = 0,
n_ctx: int = 4096,
verbose: bool = False,
chat_format: str = "chatml",
) -> None:
try:
from llama_cpp import Llama
except ImportError as e:
raise ImportError(
"llama-cpp-python is required: pip install llama-cpp-python"
) from e
if chat_format not in CHAT_FORMATS:
raise ValueError(
f"Unknown chat_format {chat_format!r}; choose from {list(CHAT_FORMATS)}"
)
self._model_path = str(Path(model_path).expanduser().resolve())
self._session_id = session_id
self._state_dir = Path(state_dir).expanduser()
self._state_dir.mkdir(parents=True, exist_ok=True)
self._chat_format = chat_format
self._llama = Llama(
model_path=self._model_path,
n_ctx=n_ctx,
n_gpu_layers=n_gpu_layers,
verbose=verbose,
)
# Auto-load existing state if present
sp = self._state_path()
if sp.exists():
core.load_state(self._llama, self._model_path, str(sp))
# ── Public methods ─────────────────────────────────────────────
def chat(self, prompt: str, max_tokens: int = 512) -> 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.
"""
template, stops = CHAT_FORMATS[self._chat_format]
wrapped = template.format(prompt=prompt)
result = self._llama(
wrapped,
max_tokens=max_tokens,
echo=False,
stop=stops,
)
return result["choices"][0]["text"].strip()
def save(self, session_id: Optional[str] = None) -> Path:
"""Persist the current state. Returns the path written."""
path = self._state_path(session_id)
core.save_state(self._llama, self._model_path, str(path))
return path
def load(self, session_id: Optional[str] = None) -> None:
"""Restore state from a (possibly different) session."""
path = self._state_path(session_id)
if not path.exists():
raise FileNotFoundError(f"No state file found: {path}")
core.load_state(self._llama, self._model_path, str(path))
@property
def state_size(self) -> int:
"""Current serialised byte size of the SSM hidden state."""
return core.get_state_size(self._llama)
@property
def session_id(self) -> str:
return self._session_id
@property
def model_path(self) -> str:
return self._model_path
# ── Private ────────────────────────────────────────────────────
def _state_path(self, session_id: Optional[str] = None) -> Path:
sid = session_id if session_id is not None else self._session_id
return self._state_dir / f"{sid}.memb"
def __repr__(self) -> str:
return (
f"Session(model={Path(self._model_path).name!r}, "
f"session_id={self._session_id!r}, "
f"state_dir={str(self._state_dir)!r})"
)