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
@@ -0,0 +1,21 @@
|
||||
"""
|
||||
memba — persistent memory layer for SSM-based LLMs.
|
||||
|
||||
Quickstart
|
||||
----------
|
||||
from memba import Session
|
||||
|
||||
s = Session("falcon-mamba-7b-Q4_K_M.gguf", session_id="research")
|
||||
print(s.chat("The Transformer architecture was introduced in 2017."))
|
||||
s.save()
|
||||
|
||||
# Later — same session, picks up where it left off
|
||||
s2 = Session("falcon-mamba-7b-Q4_K_M.gguf", session_id="research")
|
||||
print(s2.chat("Who were the authors?"))
|
||||
"""
|
||||
|
||||
from .session import Session
|
||||
from .core import save_state, load_state, get_state_size, compute_model_id
|
||||
|
||||
__all__ = ["Session", "save_state", "load_state", "get_state_size", "compute_model_id"]
|
||||
__version__ = "0.1.0"
|
||||
@@ -0,0 +1,257 @@
|
||||
"""
|
||||
memba CLI — typer-based entry point.
|
||||
|
||||
Sub-commands
|
||||
------------
|
||||
memba chat --model <gguf> --session <name> [--gpu-layers N] → REPL
|
||||
memba run --model <gguf> --prompt <text> [--save-state <f>] → one-shot
|
||||
memba list [--state-dir <dir>] → list sessions
|
||||
memba rm <session> [--state-dir <dir>] → delete session
|
||||
memba info <session> [--state-dir <dir>] → show metadata
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import struct
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import typer
|
||||
from rich.console import Console
|
||||
from rich.table import Table
|
||||
from rich.prompt import Prompt
|
||||
|
||||
from . import core
|
||||
from .session import Session
|
||||
|
||||
app = Console()
|
||||
cli = typer.Typer(
|
||||
name="memba",
|
||||
help="Persistent memory layer for SSM-based LLMs (Falcon-Mamba, Zamba).",
|
||||
add_completion=False,
|
||||
)
|
||||
console = Console()
|
||||
|
||||
_DEFAULT_STATE_DIR = "~/.memba/states"
|
||||
|
||||
# ── Helpers ────────────────────────────────────────────────────────
|
||||
|
||||
def _state_dir_path(state_dir: str) -> Path:
|
||||
return Path(state_dir).expanduser()
|
||||
|
||||
|
||||
def _list_sessions(state_dir: Path) -> list[Path]:
|
||||
if not state_dir.exists():
|
||||
return []
|
||||
return sorted(state_dir.glob("*.memb"))
|
||||
|
||||
|
||||
def _parse_memb_header(path: Path) -> dict:
|
||||
"""Return header fields from a .memb file without loading the full blob."""
|
||||
HEADER_FMT = "<4sI64sIIQ"
|
||||
HEADER_SIZE = struct.calcsize(HEADER_FMT)
|
||||
raw = path.read_bytes()
|
||||
if len(raw) < HEADER_SIZE:
|
||||
return {}
|
||||
magic, version, model_id, n_ctx, _, data_size = struct.unpack_from(HEADER_FMT, raw)
|
||||
return {
|
||||
"magic": magic,
|
||||
"version": version,
|
||||
"model_id": model_id.rstrip(b"\x00").decode("ascii", errors="replace"),
|
||||
"n_ctx": n_ctx,
|
||||
"data_size": data_size,
|
||||
"file_size": path.stat().st_size,
|
||||
}
|
||||
|
||||
|
||||
# ── chat — interactive REPL ────────────────────────────────────────
|
||||
|
||||
@cli.command()
|
||||
def chat(
|
||||
model: str = typer.Option(..., "--model", "-m", help="Path to GGUF model"),
|
||||
session: str = typer.Option("default", "--session", "-s", help="Session name"),
|
||||
gpu_layers: int = typer.Option(0, "--gpu-layers", "-g", help="GPU layers (0=CPU, -1=all)"),
|
||||
n_ctx: int = typer.Option(4096, "--n-ctx", help="Context size"),
|
||||
state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
|
||||
verbose: bool = typer.Option(False, "--verbose", "-v"),
|
||||
max_tokens: int = typer.Option(512, "--max-tokens", help="Max tokens per turn"),
|
||||
chat_format: str = typer.Option("chatml", "--chat-format",
|
||||
help="Prompt template: chatml (instruct) | raw (base)"),
|
||||
) -> None:
|
||||
"""Interactive REPL with auto-save on exit (Ctrl-C or /exit)."""
|
||||
console.print(f"[bold cyan]memba[/bold cyan] — session [green]{session!r}[/green]")
|
||||
console.print(f"Model : [dim]{model}[/dim]")
|
||||
console.print(f"GPU : {gpu_layers} layers Format: [magenta]{chat_format}[/magenta]")
|
||||
console.print("Type [bold]/save[/bold] to checkpoint, [bold]/exit[/bold] or Ctrl-C to quit.\n")
|
||||
|
||||
with console.status("Loading model…"):
|
||||
sess = Session(
|
||||
model_path=model,
|
||||
session_id=session,
|
||||
state_dir=state_dir,
|
||||
n_gpu_layers=gpu_layers,
|
||||
n_ctx=n_ctx,
|
||||
verbose=verbose,
|
||||
chat_format=chat_format,
|
||||
)
|
||||
console.print(f"[dim]State size: {sess.state_size:,} bytes[/dim]\n")
|
||||
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
user_input = Prompt.ask("[bold]You[/bold]")
|
||||
except (EOFError, KeyboardInterrupt):
|
||||
break
|
||||
|
||||
if not user_input.strip():
|
||||
continue
|
||||
if user_input.strip() == "/exit":
|
||||
break
|
||||
if user_input.strip() == "/save":
|
||||
path = sess.save()
|
||||
console.print(f"[dim]Saved → {path}[/dim]")
|
||||
continue
|
||||
|
||||
with console.status("Thinking…"):
|
||||
reply = sess.chat(user_input, max_tokens=max_tokens)
|
||||
console.print(f"[bold green]Assistant[/bold green]: {reply}\n")
|
||||
|
||||
except KeyboardInterrupt:
|
||||
pass
|
||||
|
||||
console.print("\n[dim]Saving state…[/dim]")
|
||||
path = sess.save()
|
||||
console.print(f"[bold]Session saved:[/bold] {path}")
|
||||
|
||||
|
||||
# ── run — one-shot with optional save/load ─────────────────────────
|
||||
|
||||
@cli.command()
|
||||
def run(
|
||||
model: str = typer.Option(..., "--model", "-m", help="Path to GGUF model"),
|
||||
prompt: str = typer.Option(..., "--prompt", "-p", help="Prompt text"),
|
||||
save_state: Optional[str] = typer.Option(None, "--save-state", help="Write state to this path"),
|
||||
load_state: Optional[str] = typer.Option(None, "--load-state", help="Load state from this path"),
|
||||
gpu_layers: int = typer.Option(0, "--gpu-layers", "-g"),
|
||||
n_ctx: int = typer.Option(4096, "--n-ctx"),
|
||||
max_tokens: int = typer.Option(512, "--max-tokens"),
|
||||
verbose: bool = typer.Option(False, "--verbose", "-v"),
|
||||
) -> None:
|
||||
"""One-shot generation with optional state save/load."""
|
||||
from llama_cpp import Llama
|
||||
|
||||
with console.status("Loading model…"):
|
||||
llama = Llama(
|
||||
model_path=model,
|
||||
n_ctx=n_ctx,
|
||||
n_gpu_layers=gpu_layers,
|
||||
verbose=verbose,
|
||||
)
|
||||
|
||||
if load_state:
|
||||
with console.status(f"Loading state from {load_state}…"):
|
||||
core.load_state(llama, model, load_state)
|
||||
console.print(f"[dim]State loaded: {load_state}[/dim]", file=sys.stderr)
|
||||
|
||||
result = llama(prompt, max_tokens=max_tokens, echo=False)
|
||||
print(result["choices"][0]["text"])
|
||||
|
||||
if save_state:
|
||||
core.save_state(llama, model, save_state)
|
||||
console.print(f"[dim]State saved: {save_state}[/dim]", file=sys.stderr)
|
||||
|
||||
|
||||
# ── list ───────────────────────────────────────────────────────────
|
||||
|
||||
@cli.command(name="list")
|
||||
def list_sessions(
|
||||
state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
|
||||
) -> None:
|
||||
"""List all saved sessions."""
|
||||
sdir = _state_dir_path(state_dir)
|
||||
files = _list_sessions(sdir)
|
||||
|
||||
if not files:
|
||||
console.print(f"[dim]No sessions in {sdir}[/dim]")
|
||||
return
|
||||
|
||||
table = Table(title=f"Sessions in {sdir}", show_lines=False)
|
||||
table.add_column("Session", style="bold cyan")
|
||||
table.add_column("State (B)", justify="right")
|
||||
table.add_column("File (B)", justify="right")
|
||||
table.add_column("model_id", style="dim", no_wrap=True)
|
||||
|
||||
for f in files:
|
||||
h = _parse_memb_header(f)
|
||||
table.add_row(
|
||||
f.stem,
|
||||
f"{h.get('data_size', '?'):,}" if isinstance(h.get('data_size'), int) else "?",
|
||||
f"{h.get('file_size', '?'):,}" if isinstance(h.get('file_size'), int) else "?",
|
||||
h.get("model_id", "?")[:16] + "…",
|
||||
)
|
||||
|
||||
console.print(table)
|
||||
|
||||
|
||||
# ── rm ─────────────────────────────────────────────────────────────
|
||||
|
||||
@cli.command()
|
||||
def rm(
|
||||
session: str = typer.Argument(..., help="Session name to remove"),
|
||||
state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
|
||||
yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation"),
|
||||
) -> None:
|
||||
"""Delete a saved session."""
|
||||
sdir = _state_dir_path(state_dir)
|
||||
path = sdir / f"{session}.memb"
|
||||
|
||||
if not path.exists():
|
||||
console.print(f"[red]Session not found:[/red] {path}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
if not yes:
|
||||
confirmed = typer.confirm(f"Delete {path}?")
|
||||
if not confirmed:
|
||||
raise typer.Abort()
|
||||
|
||||
path.unlink()
|
||||
console.print(f"[dim]Deleted:[/dim] {path}")
|
||||
|
||||
|
||||
# ── info ───────────────────────────────────────────────────────────
|
||||
|
||||
@cli.command()
|
||||
def info(
|
||||
session: str = typer.Argument(..., help="Session name"),
|
||||
state_dir: str = typer.Option(_DEFAULT_STATE_DIR, "--state-dir"),
|
||||
) -> None:
|
||||
"""Show metadata stored in a session's state file."""
|
||||
sdir = _state_dir_path(state_dir)
|
||||
path = sdir / f"{session}.memb"
|
||||
|
||||
if not path.exists():
|
||||
console.print(f"[red]Session not found:[/red] {path}")
|
||||
raise typer.Exit(1)
|
||||
|
||||
h = _parse_memb_header(path)
|
||||
if not h:
|
||||
console.print("[red]Cannot parse header (truncated file?)[/red]")
|
||||
raise typer.Exit(1)
|
||||
|
||||
console.print(f"[bold]Session:[/bold] {session}")
|
||||
console.print(f" File : {path}")
|
||||
console.print(f" Magic : {h['magic']!r}")
|
||||
console.print(f" Version : {h['version']}")
|
||||
console.print(f" model_id: {h['model_id']}")
|
||||
console.print(f" n_ctx : {h['n_ctx']}")
|
||||
console.print(f" State : {h['data_size']:,} bytes")
|
||||
console.print(f" File : {h['file_size']:,} bytes")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
cli()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,189 @@
|
||||
"""
|
||||
Low-level state I/O — reads/writes the MEMB file format using llama-cpp-python's
|
||||
exposed C functions directly. No ABI conflict: we reuse the libllama.so that
|
||||
llama-cpp-python already loaded instead of linking our own copy.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import ctypes
|
||||
import hashlib
|
||||
import struct
|
||||
import zlib
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from llama_cpp import Llama
|
||||
|
||||
# ── File format ───────────────────────────────────────────────────
|
||||
MAGIC = b"MEMB"
|
||||
VERSION = 1
|
||||
MODEL_ID_LEN = 64
|
||||
|
||||
# little-endian: 4s magic | I version | 64s model_id | I n_ctx | I llama_ver | Q data_size
|
||||
_HEADER_FMT = "<4sI64sIIQ"
|
||||
_HEADER_SIZE = struct.calcsize(_HEADER_FMT) # == 88
|
||||
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────
|
||||
|
||||
def compute_model_id(model_path: str) -> bytes:
|
||||
"""SHA-256 of first 1 KiB of the GGUF file, hex-encoded and zero-padded to 64 bytes."""
|
||||
h = hashlib.sha256()
|
||||
try:
|
||||
with open(model_path, "rb") as f:
|
||||
h.update(f.read(1024))
|
||||
except OSError:
|
||||
h.update(model_path.encode("utf-8", errors="replace"))
|
||||
digest = h.hexdigest().encode("ascii") # 64 chars exactly for SHA-256 hex
|
||||
return digest.ljust(MODEL_ID_LEN, b"\x00")[:MODEL_ID_LEN]
|
||||
|
||||
|
||||
def _lib_and_ctx(llama_model: "Llama"):
|
||||
"""Return (llama_cpp low-level module, raw llama_context_p pointer).
|
||||
|
||||
Modern llama-cpp-python (≥0.3) wraps the raw `llama_context*` in a
|
||||
`_LlamaContext` helper object — the actual ctypes pointer lives one
|
||||
attribute deeper. Older versions exposed the pointer directly.
|
||||
"""
|
||||
try:
|
||||
from llama_cpp import llama_cpp as lib
|
||||
except ImportError as e:
|
||||
raise ImportError("llama-cpp-python is not installed: pip install llama-cpp-python") from e
|
||||
|
||||
wrapper = getattr(llama_model, "_ctx", None) or getattr(llama_model, "ctx", None)
|
||||
if wrapper is None:
|
||||
raise RuntimeError(
|
||||
"Cannot find _ctx attribute on Llama instance — unsupported llama-cpp-python version"
|
||||
)
|
||||
|
||||
# If wrapper is already a ctypes pointer (older llama-cpp-python), use directly
|
||||
if isinstance(wrapper, (int, ctypes.c_void_p)) or hasattr(wrapper, "_type_"):
|
||||
return lib, wrapper
|
||||
|
||||
# Otherwise unwrap one level: `_LlamaContext.ctx` holds the raw pointer
|
||||
# (may be an int address, ctypes pointer, or c_void_p depending on version)
|
||||
for attr in ("ctx", "context", "_ctx"):
|
||||
raw = getattr(wrapper, attr, None)
|
||||
if raw is not None and (isinstance(raw, (int, ctypes.c_void_p)) or hasattr(raw, "_type_")):
|
||||
return lib, raw
|
||||
|
||||
raise RuntimeError(
|
||||
f"Could not extract raw llama_context from {type(wrapper).__name__}; "
|
||||
"your llama-cpp-python version may have changed its internals."
|
||||
)
|
||||
|
||||
|
||||
def _state_get_size(lib, ctx) -> int:
|
||||
try:
|
||||
return lib.llama_state_get_size(ctx) # post-2024 API
|
||||
except AttributeError:
|
||||
return lib.llama_get_state_size(ctx) # VERIFY: older API fallback
|
||||
|
||||
|
||||
def _state_get_data(lib, ctx, buf: ctypes.Array, size: int) -> int:
|
||||
try:
|
||||
return lib.llama_state_get_data(ctx, buf, size) # post-2024 API
|
||||
except AttributeError:
|
||||
return lib.llama_copy_state_data(ctx, buf) # VERIFY: older API fallback
|
||||
|
||||
|
||||
def _state_set_data(lib, ctx, buf: ctypes.Array, size: int) -> int:
|
||||
try:
|
||||
return lib.llama_state_set_data(ctx, buf, size) # post-2024 API
|
||||
except AttributeError:
|
||||
return lib.llama_set_state_data(ctx, buf) # VERIFY: older API fallback
|
||||
|
||||
|
||||
def _n_ctx(lib, ctx) -> int:
|
||||
try:
|
||||
return lib.llama_n_ctx(ctx)
|
||||
except AttributeError:
|
||||
return 0
|
||||
|
||||
|
||||
# ── Public API ────────────────────────────────────────────────────
|
||||
|
||||
def get_state_size(llama_model: "Llama") -> int:
|
||||
"""Return the byte size of the current SSM hidden state."""
|
||||
lib, ctx = _lib_and_ctx(llama_model)
|
||||
return _state_get_size(lib, ctx)
|
||||
|
||||
|
||||
def save_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
|
||||
"""Serialise the current SSM state to *file_path* in MEMB format."""
|
||||
lib, ctx = _lib_and_ctx(llama_model)
|
||||
|
||||
size = _state_get_size(lib, ctx)
|
||||
if size == 0:
|
||||
raise RuntimeError("llama context returned state size 0 — is it initialised?")
|
||||
|
||||
buf = (ctypes.c_uint8 * size)()
|
||||
written = _state_get_data(lib, ctx, buf, size)
|
||||
if written == 0:
|
||||
raise RuntimeError("llama_state_get_data returned 0 bytes")
|
||||
|
||||
data = bytes(buf[:written])
|
||||
crc = zlib.crc32(data) & 0xFFFFFFFF
|
||||
|
||||
model_id = compute_model_id(model_path)
|
||||
n_ctx = _n_ctx(lib, ctx)
|
||||
|
||||
header = struct.pack(
|
||||
_HEADER_FMT,
|
||||
MAGIC,
|
||||
VERSION,
|
||||
model_id,
|
||||
n_ctx,
|
||||
0, # llama_ver — reserved
|
||||
len(data),
|
||||
)
|
||||
|
||||
out = Path(file_path)
|
||||
out.parent.mkdir(parents=True, exist_ok=True)
|
||||
out.write_bytes(header + data + struct.pack("<I", crc))
|
||||
|
||||
|
||||
def load_state(llama_model: "Llama", model_path: str, file_path: str) -> None:
|
||||
"""Restore SSM state from *file_path* into *llama_model*."""
|
||||
lib, ctx = _lib_and_ctx(llama_model)
|
||||
expected_model_id = compute_model_id(model_path)
|
||||
|
||||
raw = Path(file_path).read_bytes()
|
||||
if len(raw) < _HEADER_SIZE + 4:
|
||||
raise ValueError(f"State file too small: {file_path}")
|
||||
|
||||
magic, version, file_model_id, n_ctx, llama_ver, data_size = struct.unpack_from(
|
||||
_HEADER_FMT, raw, 0
|
||||
)
|
||||
|
||||
if magic != MAGIC:
|
||||
raise ValueError(f"Bad magic: expected {MAGIC!r}, got {magic!r}")
|
||||
if version != VERSION:
|
||||
raise ValueError(f"Unsupported state version {version} (expected {VERSION})")
|
||||
if file_model_id != expected_model_id:
|
||||
raise ValueError(
|
||||
"Model identity mismatch — state file was created with a different GGUF.\n"
|
||||
f" file model_id : {file_model_id.rstrip(b'\\x00').decode()}\n"
|
||||
f" current model : {expected_model_id.rstrip(b'\\x00').decode()}"
|
||||
)
|
||||
|
||||
offset = _HEADER_SIZE
|
||||
if len(raw) < offset + data_size + 4:
|
||||
raise ValueError("Truncated state file")
|
||||
|
||||
data = raw[offset : offset + data_size]
|
||||
stored_crc = struct.unpack_from("<I", raw, offset + data_size)[0]
|
||||
computed = zlib.crc32(data) & 0xFFFFFFFF
|
||||
|
||||
if computed != stored_crc:
|
||||
raise ValueError(
|
||||
f"CRC-32 mismatch (stored={stored_crc:#010x}, computed={computed:#010x}) — "
|
||||
"file may be corrupted"
|
||||
)
|
||||
|
||||
buf = (ctypes.c_uint8 * len(data)).from_buffer_copy(data)
|
||||
restored = _state_set_data(lib, ctx, buf, len(data))
|
||||
if restored == 0:
|
||||
raise RuntimeError("llama_state_set_data returned 0 — state restore failed")
|
||||
@@ -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