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:
@@ -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})"
|
||||
)
|
||||
Reference in New Issue
Block a user