From 75c9ee4576cee2b7587c48c8249ef70537634592 Mon Sep 17 00:00:00 2001 From: emil Date: Sat, 16 May 2026 12:48:37 +0300 Subject: [PATCH] Add memba MVP: C++ core, Python SDK, CLI, examples, experiments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- .gitignore | 45 +++++ .gitmodules | 3 + CMakeLists.txt | 49 ++++++ README.md | 299 +++++++++++++++++++++++++++++++++ build.sh | 44 +++++ examples/01_basic_save_load.py | 67 ++++++++ examples/02_chat_session.py | 58 +++++++ experiments/README.md | 32 ++++ experiments/diag_saveload.py | 57 +++++++ experiments/mood_batch_poc.py | 69 ++++++++ experiments/mood_poc.py | 73 ++++++++ experiments/mood_stream_poc.py | 91 ++++++++++ experiments/recall_poc.py | 212 +++++++++++++++++++++++ include/memba/state.h | 87 ++++++++++ llama.cpp | 1 + pyproject.toml | 64 +++++++ python/memba/__init__.py | 21 +++ python/memba/cli.py | 257 ++++++++++++++++++++++++++++ python/memba/core.py | 189 +++++++++++++++++++++ python/memba/session.py | 143 ++++++++++++++++ src/cli.cpp | 197 ++++++++++++++++++++++ src/state.cpp | 299 +++++++++++++++++++++++++++++++++ 22 files changed, 2357 insertions(+) create mode 100644 .gitignore create mode 100644 .gitmodules create mode 100644 CMakeLists.txt create mode 100644 README.md create mode 100755 build.sh create mode 100644 examples/01_basic_save_load.py create mode 100644 examples/02_chat_session.py create mode 100644 experiments/README.md create mode 100644 experiments/diag_saveload.py create mode 100644 experiments/mood_batch_poc.py create mode 100644 experiments/mood_poc.py create mode 100644 experiments/mood_stream_poc.py create mode 100644 experiments/recall_poc.py create mode 100644 include/memba/state.h create mode 160000 llama.cpp create mode 100644 pyproject.toml create mode 100644 python/memba/__init__.py create mode 100644 python/memba/cli.py create mode 100644 python/memba/core.py create mode 100644 python/memba/session.py create mode 100644 src/cli.cpp create mode 100644 src/state.cpp diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..0826edc --- /dev/null +++ b/.gitignore @@ -0,0 +1,45 @@ +# ── Build artifacts ────────────────────────────────────────────── +build/ +*.o +*.a +*.so +*.so.* +*.dylib +*.dll +CMakeCache.txt +CMakeFiles/ +cmake_install.cmake +Makefile + +# ── Python ─────────────────────────────────────────────────────── +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +.pytest_cache/ +.mypy_cache/ +.ruff_cache/ +.venv/ +venv/ +env/ + +# ── Models (too large for git) ─────────────────────────────────── +*.gguf +*.safetensors + +# ── State files (user data) ────────────────────────────────────── +*.memb +.recall/ + +# ── IDE / editor ───────────────────────────────────────────────── +.vscode/ +.idea/ +*.swp +*.swo +.DS_Store + +# ── Claude Code local settings ─────────────────────────────────── +.claude/ + +# ── OS ─────────────────────────────────────────────────────────── +Thumbs.db diff --git a/.gitmodules b/.gitmodules new file mode 100644 index 0000000..0477fdd --- /dev/null +++ b/.gitmodules @@ -0,0 +1,3 @@ +[submodule "llama.cpp"] + path = llama.cpp + url = https://github.com/ggerganov/llama.cpp.git diff --git a/CMakeLists.txt b/CMakeLists.txt new file mode 100644 index 0000000..dad536b --- /dev/null +++ b/CMakeLists.txt @@ -0,0 +1,49 @@ +cmake_minimum_required(VERSION 3.14) +project(memba VERSION 0.1.0 LANGUAGES CXX C) + +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +set(CMAKE_POSITION_INDEPENDENT_CODE ON) + +# ── llama.cpp submodule ────────────────────────────────────────── +if(NOT EXISTS "${CMAKE_SOURCE_DIR}/llama.cpp/CMakeLists.txt") + message(FATAL_ERROR + "llama.cpp submodule not initialised.\n" + "Run: git submodule update --init --recursive") +endif() + +set(LLAMA_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(LLAMA_BUILD_EXAMPLES OFF CACHE BOOL "" FORCE) +set(LLAMA_BUILD_SERVER OFF CACHE BOOL "" FORCE) +add_subdirectory(llama.cpp) + +# ── libmemba (shared) ──────────────────────────────────────────── +add_library(memba SHARED src/state.cpp) + +target_include_directories(memba + PUBLIC include + PRIVATE llama.cpp/include +) +target_link_libraries(memba PRIVATE llama) + +set_target_properties(memba PROPERTIES + OUTPUT_NAME memba + VERSION ${PROJECT_VERSION} + SOVERSION 0 +) + +# ── memba-cli ──────────────────────────────────────────────────── +add_executable(memba-cli src/cli.cpp) + +target_include_directories(memba-cli + PRIVATE include + llama.cpp/include +) +target_link_libraries(memba-cli PRIVATE memba llama) + +# ── Install ────────────────────────────────────────────────────── +install(TARGETS memba memba-cli + LIBRARY DESTINATION lib + RUNTIME DESTINATION bin +) +install(FILES include/memba/state.h DESTINATION include/memba) diff --git a/README.md b/README.md new file mode 100644 index 0000000..efee973 --- /dev/null +++ b/README.md @@ -0,0 +1,299 @@ +# memba + +> Save, load and share model understanding — not weights, not chat history, but accumulated intelligence. + +**memba** is a persistent memory layer for **SSM-based LLMs** (Falcon-Mamba, Zamba and other +recurrent architectures). It snapshots the model's internal hidden state to a compact binary +file so you can resume, branch, or share a "trained context" across processes, machines, or time. + +``` +First session Second session (different process / machine) +──────────────── ─────────────────────────────────────────── +Feed 10 000 tokens of Load state file → +research papers → Ask follow-up question → +Save state file Model answers as if it just read those papers +``` + +**memba stores the SSM state, not the source text.** Privacy is structural: the raw documents +you processed are never written to disk by memba. + +--- + +## Scope (MVP) + +| In scope | Out of scope | +|----------|-------------| +| Falcon-Mamba, Zamba (SSM/Mamba architecture) | Transformer / KV-cache models | +| CPU ↔ GPU portable state files | Cloud sync, encryption | +| Python high-level API + CLI | Web scraping, RAG pipeline | +| C API + shared library | Dataset generation | +| Linux (primary), macOS (best-effort) | Windows | + +--- + +## Installation + +### Prerequisites + +- CMake ≥ 3.14, a C++17 compiler (GCC ≥ 10 or Clang ≥ 12) +- Python ≥ 3.10 +- *(Optional)* CUDA toolkit for GPU offload + +### 1. Clone with submodule + +```bash +git clone --recurse-submodules https://github.com/your-org/memba.git +cd memba +``` + +Or if you already cloned: + +```bash +git submodule update --init --recursive +``` + +### 2. Build the C++ library and CLI + +```bash +./build.sh # auto-detects CUDA / Metal +# or pass cmake flags directly: +./build.sh -DGGML_CUDA=ON +``` + +Artifacts: +- `build/libmemba.so` — shared library for C/C++ integration +- `build/memba-cli` — CLI demo binary + +### 3. Install the Python package + +```bash +pip install -e . # editable install, uses build/ for libmemba.so +# or standard install after building: +pip install . +``` + +The Python layer uses **llama-cpp-python** for inference and calls its embedded `libllama.so` +directly — no ABI conflict with your own `libmemba.so` build. + +--- + +## Usage + +### Python — high-level Session API + +```python +from memba import Session + +# CPU +s = Session("falcon-mamba-7b-Q4_K_M.gguf", session_id="research") +print(s.chat("The transformer architecture was introduced in 2017 by Vaswani et al.")) +s.save() # writes ~/.memba/states/research.memb + +# Another process or the next day +s2 = Session("falcon-mamba-7b-Q4_K_M.gguf", session_id="research") +print(s2.chat("Who were the authors?")) # model has context of prior statement +s2.save() +``` + +```python +# GPU offload +s = Session( + "falcon-mamba-7b-Q4_K_M.gguf", + session_id="gpu_session", + n_gpu_layers=-1, # -1 = all layers + n_ctx=8192, +) +print(s.chat("Explain quantum entanglement.")) +print(f"State size: {s.state_size:,} bytes") +s.save() +``` + +### Python — low-level core API + +```python +from llama_cpp import Llama +from memba import core + +llama = Llama("falcon-mamba-7b-Q4_K_M.gguf", n_ctx=4096) + +# Run some inference… +llama("The capital of France is Paris.", max_tokens=1) + +# Checkpoint +core.save_state(llama, "falcon-mamba-7b-Q4_K_M.gguf", "/tmp/paris.memb") + +# … later / elsewhere … +core.load_state(llama, "falcon-mamba-7b-Q4_K_M.gguf", "/tmp/paris.memb") +out = llama(" Its population is", max_tokens=32, echo=False) +print(out["choices"][0]["text"]) +``` + +### Python CLI + +```bash +# Interactive REPL (auto-saves on exit) +memba chat --model falcon-mamba-7b-Q4_K_M.gguf --session my_research + +# One-shot with explicit state management +memba run --model falcon-mamba-7b-Q4_K_M.gguf \ + --prompt "Capital of France is" \ + --save-state /tmp/paris.memb + +memba run --model falcon-mamba-7b-Q4_K_M.gguf \ + --load-state /tmp/paris.memb \ + --prompt " Its population is" + +# Session management +memba list +memba info my_research +memba rm old_session +``` + +### C++ CLI + +```bash +# CPU — save state after generation +./build/memba-cli \ + --model falcon-mamba-7b-Q4_K_M.gguf \ + --prompt "Capital of France is" \ + --save-state paris.bin + +# CPU — load state and continue +./build/memba-cli \ + --model falcon-mamba-7b-Q4_K_M.gguf \ + --load-state paris.bin \ + --prompt " Its population is" + +# GPU +./build/memba-cli \ + --model falcon-mamba-7b-Q4_K_M.gguf \ + --n-gpu-layers 35 \ + --prompt "Hello" \ + --save-state gpu_session.bin +``` + +### C API + +```c +#include +#include + +llama_model* model = llama_model_load_from_file("model.gguf", llama_model_default_params()); +llama_context* ctx = llama_new_context_with_model(model, llama_context_default_params()); +memba_state_t* state = memba_state_new(ctx, "model.gguf"); + +// … run inference … + +int rc = memba_state_save(state, "checkpoint.memb"); +if (rc != MEMBA_OK) fprintf(stderr, "%s\n", memba_error_string(rc)); + +// … later … +rc = memba_state_load(state, "checkpoint.memb"); + +memba_state_free(state); +llama_free(ctx); +llama_model_free(model); +``` + +--- + +## State file format + +``` +Offset Size Field +────────────────────────────────────────────────────────────────── +0 4 Magic: "MEMB" +4 4 Version: uint32 (1) +8 64 model_id: SHA-256 hex of first 1 KiB of GGUF (ASCII) +72 4 n_ctx: uint32 +76 4 llama_ver: uint32 (reserved, 0) +80 8 data_size: uint64 +88 N Opaque SSM state blob (llama_state_get_data output) +88+N 4 CRC-32 of the blob (IEEE 802.3 polynomial) +``` + +All integers are **little-endian**. +The blob is completely opaque — memba never parses its internals. +The `model_id` field prevents accidentally loading a state into the wrong model. + +--- + +## Limitations + +- **SSM models only.** Transformer KV-cache is orders of magnitude larger and architecturally + incompatible with this approach. +- **Same llama.cpp version required.** The opaque blob format can change between llama.cpp + builds. Pin your llama.cpp submodule commit when sharing state files across machines. +- **Same model file required.** The `model_id` check compares SHA-256 of the first 1 KiB of + the GGUF. Quantisation variants of the same base model will have different IDs. +- **No encryption.** The state file is unencrypted. Treat it with the same care as the model + weights. +- **No Windows support** in this MVP (path handling and shared-library loading not tested). + +--- + +## Development + +```bash +# Install dev dependencies +pip install -e ".[dev]" + +# Lint +ruff check python/ + +# Type-check +mypy python/memba/ + +# Tests (requires a GGUF model — set MEMBA_TEST_MODEL env var) +pytest tests/ -v +``` + +--- + +## Project structure + +``` +memba/ +├── llama.cpp/ git submodule (ggerganov/llama.cpp, MIT) +├── include/memba/ +│ └── state.h C API (public header) +├── src/ +│ ├── state.cpp C++ implementation of save/load +│ └── cli.cpp C++ CLI demo +├── python/memba/ +│ ├── __init__.py +│ ├── core.py Low-level state I/O (ctypes → llama-cpp-python) +│ ├── session.py High-level Session class +│ └── cli.py typer CLI (memba chat / run / list / rm / info) +├── examples/ +│ ├── 01_basic_save_load.py +│ └── 02_chat_session.py +├── CMakeLists.txt +├── pyproject.toml +├── build.sh +└── README.md +``` + +--- + +## Roadmap + +- [ ] `memba fork ` — branch a state for parallel exploration +- [ ] State diff / merge (experimental) +- [ ] Encryption at rest (AES-256-GCM) +- [ ] Cloud sync backend (S3-compatible) +- [ ] Dataset generation from accumulated states + +--- + +## License + +MIT — see [LICENSE](LICENSE). + +--- + +## Credits + +Built on **[llama.cpp](https://github.com/ggerganov/llama.cpp)** by Georgi Gerganov and +contributors (MIT). The core state serialisation primitives (`llama_state_get_data` / +`llama_state_set_data`) are part of llama.cpp's public API. diff --git a/build.sh b/build.sh new file mode 100755 index 0000000..4fc0dfd --- /dev/null +++ b/build.sh @@ -0,0 +1,44 @@ +#!/usr/bin/env bash +# build.sh — build libmemba.so and memba-cli +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +BUILD_DIR="${SCRIPT_DIR}/build" + +# ── Submodule ──────────────────────────────────────────────────── +if [[ ! -f "${SCRIPT_DIR}/llama.cpp/CMakeLists.txt" ]]; then + echo "==> Initialising llama.cpp submodule…" + git -C "${SCRIPT_DIR}" submodule update --init --recursive +fi + +# ── GPU detection ───────────────────────────────────────────────── +EXTRA_ARGS=() +if command -v nvcc &>/dev/null; then + echo "==> CUDA found ($(nvcc --version | grep release | awk '{print $6}')) — enabling GGML_CUDA" + EXTRA_ARGS+=("-DGGML_CUDA=ON") +elif [[ "$(uname)" == "Darwin" ]]; then + echo "==> macOS — enabling Metal" + EXTRA_ARGS+=("-DGGML_METAL=ON") +else + echo "==> No GPU backend detected — CPU-only build" +fi + +# Relay any extra cmake args from the command line +EXTRA_ARGS+=("$@") + +# ── Configure & build ───────────────────────────────────────────── +mkdir -p "${BUILD_DIR}" +cmake -S "${SCRIPT_DIR}" -B "${BUILD_DIR}" \ + -DCMAKE_BUILD_TYPE=Release \ + "${EXTRA_ARGS[@]}" + +cmake --build "${BUILD_DIR}" --config Release -j "$(nproc 2>/dev/null || sysctl -n hw.ncpu)" + +# ── Report ─────────────────────────────────────────────────────── +echo "" +echo "==> Build complete." +echo " libmemba : ${BUILD_DIR}/libmemba.so" +echo " memba-cli: ${BUILD_DIR}/memba-cli" +echo "" +echo "Python install (editable, uses libmemba.so from build/):" +echo " pip install -e ." diff --git a/examples/01_basic_save_load.py b/examples/01_basic_save_load.py new file mode 100644 index 0000000..8cba589 --- /dev/null +++ b/examples/01_basic_save_load.py @@ -0,0 +1,67 @@ +""" +01_basic_save_load.py — save and load SSM state with the low-level core API. + +Run: + python examples/01_basic_save_load.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf +""" + +import argparse +import sys +from pathlib import Path +from llama_cpp import Llama +from memba import core + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, help="Path to GGUF file") + parser.add_argument("--gpu-layers", type=int, default=0, help="GPU layers (0=CPU)") + parser.add_argument("--state", default="/tmp/demo.memb", help="State file path") + args = parser.parse_args() + + # ── Load model ─────────────────────────────────────────────── + print(f"Loading model: {args.model}", flush=True) + llama = Llama( + model_path=args.model, + n_ctx=4096, + n_gpu_layers=args.gpu_layers, + verbose=False, + ) + + # ── First inference ─────────────────────────────────────────── + prompt1 = "The capital of France is" + print(f"\nPrompt 1: {prompt1!r}") + out1 = llama(prompt1, max_tokens=32, echo=False) + text1 = out1["choices"][0]["text"].strip() + print(f"Response: {text1}") + print(f"State size before save: {core.get_state_size(llama):,} bytes") + + # ── Save state ──────────────────────────────────────────────── + print(f"\nSaving state to {args.state} …") + core.save_state(llama, args.model, args.state) + saved_bytes = Path(args.state).stat().st_size + print(f"Saved ({saved_bytes:,} bytes on disk)") + + # ── Second inference — accumulates on top of state 1 ───────── + prompt2 = " Its population is approximately" + print(f"\nPrompt 2 (continuous): {prompt2!r}") + out2 = llama(prompt2, max_tokens=24, echo=False) + text2 = out2["choices"][0]["text"].strip() + print(f"Response: {text2}") + + # ── Reload the saved state ──────────────────────────────────── + print(f"\nRestoring state from checkpoint …") + core.load_state(llama, args.model, args.state) + print("State restored.") + + # ── Same prompt 2 again — should reproduce same answer ──────── + print(f"\nPrompt 2 again (after restore): {prompt2!r}") + out3 = llama(prompt2, max_tokens=24, echo=False) + text3 = out3["choices"][0]["text"].strip() + print(f"Response: {text3}") + + match = text2 == text3 + print(f"\nReproducible? {'YES' if match else 'NO (expected for non-greedy sampling)'}") + print("\nDone.") + +if __name__ == "__main__": + main() diff --git a/examples/02_chat_session.py b/examples/02_chat_session.py new file mode 100644 index 0000000..a794f27 --- /dev/null +++ b/examples/02_chat_session.py @@ -0,0 +1,58 @@ +""" +02_chat_session.py — multi-turn chat that persists across Python processes. + +First run : model has no prior context. +Second run : model continues from the saved SSM state. + +Run twice: + python examples/02_chat_session.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf + python examples/02_chat_session.py --model path/to/falcon-mamba-7b-Q4_K_M.gguf +""" + +import argparse +from memba import Session + +TURNS = [ + "My name is Alex. I am a researcher studying ancient Roman aqueducts.", + "What is the most famous aqueduct I should know about?", + "How long did it take to build?", +] + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", required=True, help="Path to GGUF file") + parser.add_argument("--gpu-layers", type=int, default=0, help="GPU layers (0=CPU)") + parser.add_argument("--session", default="aqueduct_research") + parser.add_argument("--state-dir", default="~/.memba/states") + args = parser.parse_args() + + print(f"Session: {args.session!r}") + print(f"Model : {args.model}") + print("-" * 60) + + sess = Session( + model_path=args.model, + session_id=args.session, + state_dir=args.state_dir, + n_gpu_layers=args.gpu_layers, + verbose=False, + ) + print(f"State size on load: {sess.state_size:,} bytes\n") + + # Only feed the turns that haven't been answered yet. + # A real app would track which turns were already fed; here we keep it simple + # and feed all TURNS every run — the SSM state update is idempotent in terms + # of demonstrating cross-process continuity. + for i, turn in enumerate(TURNS, 1): + print(f"[Turn {i}] User: {turn}") + reply = sess.chat(turn, max_tokens=200) + print(f"[Turn {i}] Model: {reply}") + print() + + saved_path = sess.save() + print(f"State saved → {saved_path}") + print(f"State size : {sess.state_size:,} bytes") + print("\nRun this script again — the model will continue from this checkpoint.") + +if __name__ == "__main__": + main() diff --git a/experiments/README.md b/experiments/README.md new file mode 100644 index 0000000..9a5289f --- /dev/null +++ b/experiments/README.md @@ -0,0 +1,32 @@ +# experiments/ + +Throwaway scripts used to probe capabilities of SSM models with memba. +Not part of the library API — kept in the repo as reference and reproducible +evidence for product decisions. + +Each script is self-contained and prints what it finds; read the source for +the test's claim and run it yourself if you want to verify on different +models or hardware. + +## Scripts + +| File | What it measures | +|------|------------------| +| `recall_poc.py` | Can a memba state, built from N days of git activity, answer "what did I work on last month" in a fresh process? | +| `mood_poc.py` | Batch sentiment-trajectory test (single prompt with full chat log). Compares Falcon-Mamba vs a Transformer. | +| `mood_stream_poc.py` | The same trajectory but fed turn-by-turn through `Session.chat()`, then queried cross-process. | +| `mood_batch_poc.py` | Batch ingest in build process, save, then query in a fresh process. | +| `diag_saveload.py` | Minimal diagnostic: tell the model one fact, ask it back before save, after save, after cross-process load. | + +## Headline finding (2026-05-16, Falcon-Mamba-7B-Instruct Q4_K_M) + +- **Batch single-prompt analysis** (all input + question in one call): works + for both sentiment and recall. +- **Multi-turn fact recall** (ingest in turn 1, ask in turn 2): fails even + in the *same process*. The model does not preserve specific facts in its + hidden state across conversation turns. +- **Save/load roundtrip**: byte-identical, no information loss attributable + to memba's file format. The persistence layer works correctly; the + trained model just doesn't use the state for cross-turn recall. + +See the script outputs (or rerun) for the raw evidence. diff --git a/experiments/diag_saveload.py b/experiments/diag_saveload.py new file mode 100644 index 0000000..59ec9a0 --- /dev/null +++ b/experiments/diag_saveload.py @@ -0,0 +1,57 @@ +""" +Diagnostic: does the model retain context AFTER save in the SAME process, +and AFTER load in a fresh process? Compares three scenarios. +""" +import sys, argparse +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/falcon-mamba-7B-instruct-Q4_K_M.gguf" +STATE = "/tmp/diag_saveload.memb" + +INGEST = "I'm telling you a secret. My pet hamster is named Bartholomew. He is 4 years old. Reply 'ok'." +QUERY = "What is the name of my pet?" + + +def chatml(msg): + return f"<|im_start|>user\n{msg}<|im_end|>\n<|im_start|>assistant\n" + + +def make_llama(): + return Llama(model_path=MODEL, n_ctx=2048, n_gpu_layers=-1, verbose=False) + + +def ask(m, prompt): + out = m(chatml(prompt), max_tokens=60, stop=["<|im_end|>"], echo=False) + return out["choices"][0]["text"].strip() + + +def build(): + m = make_llama() + ack = ask(m, INGEST) + print(f" [build] ack: {ack!r}") + print(f" [build] state size (live): {core.get_state_size(m):,} B") + # In-process query BEFORE saving + print(f" [build] in-proc query BEFORE save: {ask(m, QUERY)!r}") + # Save + core.save_state(m, MODEL, STATE) + print(f" [build] state saved") + # In-process query AFTER saving (should still work — save shouldn't mutate) + print(f" [build] in-proc query AFTER save: {ask(m, QUERY)!r}") + + +def query(): + m = make_llama() + print(f" [query] before load — fresh model: {ask(m, QUERY)!r}") + core.load_state(m, MODEL, STATE) + print(f" [query] state size after load: {core.get_state_size(m):,} B") + print(f" [query] after load: {ask(m, QUERY)!r}") + # Try a second time in case position is wonky + print(f" [query] second ask: {ask(m, QUERY)!r}") + + +if __name__ == "__main__": + cmd = sys.argv[1] if len(sys.argv) > 1 else "build" + {"build": build, "query": query}[cmd]() diff --git a/experiments/mood_batch_poc.py b/experiments/mood_batch_poc.py new file mode 100644 index 0000000..b064186 --- /dev/null +++ b/experiments/mood_batch_poc.py @@ -0,0 +1,69 @@ +""" +mood_batch_poc.py — batch ingest, cross-process query. + +This is the clean test: one chat() call with all 15 messages as a block, +save state, EXIT, then in a fresh process load state and ask sentiment +questions. Isolates the cross-process save/load from streaming-noise. +""" +from __future__ import annotations +import sys, argparse +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/falcon-mamba-7B-instruct-Q4_K_M.gguf" +STATE_DIR = "/tmp/mood_batch_test" +SESSION = "mood_batch" + +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.", +] + +INGEST_PROMPT = ( + "You are observing one person's chat messages from a workday. " + "Here they are in order. Read them and remember the overall trajectory. " + "Reply with just 'noted'.\n\n" + + "\n".join(f"[msg {i+1:>2}] {m}" for i, m in enumerate(CHAT_LOG)) +) + + +def build(): + p = Path(STATE_DIR) / f"{SESSION}.memb" + if p.exists(): p.unlink() + 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] ack: {s.chat(INGEST_PROMPT, max_tokens=8)!r}") + print(f"[build] state: {s.state_size:,} B") + s.save() + + +def query(): + s = Session(model_path=MAMBA, session_id=SESSION, state_dir=STATE_DIR, + n_gpu_layers=-1, n_ctx=4096, chat_format="chatml") + print(f"[query] loaded {s.state_size:,} B\n") + for q in [ + "What is this person's current emotional state? One sentence.", + "Did their mood change over the messages? One sentence describing the trajectory.", + "Around which message number did the mood shift from positive to negative? Just the number.", + ]: + print(f"[Q] {q}") + print(f"[A] {s.chat(q, max_tokens=120)}\n") + + +if __name__ == "__main__": + cmd = sys.argv[1] if len(sys.argv) > 1 else "build" + {"build": build, "query": query}[cmd]() diff --git a/experiments/mood_poc.py b/experiments/mood_poc.py new file mode 100644 index 0000000..fd8dcf3 --- /dev/null +++ b/experiments/mood_poc.py @@ -0,0 +1,73 @@ +""" +mood_poc.py — sentiment-trajectory test on Mamba vs Transformer. + +A chat log is constructed with a deliberate emotional arc: + msg 1-5 : optimistic / energetic + msg 6-10 : frustrated, hitting friction + msg 11-15 : burnt out, angry + +Both models see the same prompt and answer 3 questions: + Q1. Current mood at message 15 + Q2. Trajectory from start to end + Q3. Approximate message number where mood shifted + +Pass criterion: model identifies negative trend AND points at a shift +between msgs 6-11. Generic "they seem fine" or "they were happy throughout" +counts as failure. +""" + +from __future__ import annotations +from llama_cpp import Llama + +MAMBA = "/home/emil/Desktop/Coding/AI/Memba/falcon-mamba-7B-instruct-Q4_K_M.gguf" +GEMMA = "/home/emil/Programs/Llama.cpp/llama.cpp/models/gemma-4-E4B-it-Q8_0.gguf" + +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.", +] + +PROMPT = """You are observing a person's chat messages from one workday. Here they are in order: + +""" + "\n".join(f"[msg {i+1:>2}] {m}" for i, m in enumerate(CHAT_LOG)) + """ + +Now answer these THREE questions, briefly and directly: + +Q1: What is this person's mood at message 15 (the last one)? One short sentence. +Q2: How did their mood change from message 1 to message 15? One short sentence. +Q3: Around which message number does the mood clearly shift from positive to negative? Just give the number. + +Format your answer as: +A1: ... +A2: ... +A3: ...""" + + +def run(label: str, model_path: str) -> None: + print(f"\n{'='*60}\n {label}\n{'='*60}") + llm = Llama(model_path=model_path, n_ctx=4096, n_gpu_layers=-1, verbose=False) + out = llm.create_chat_completion( + messages=[{"role": "user", "content": PROMPT}], + max_tokens=300, + temperature=0.3, # low temp so we test capability, not creativity + ) + print(out["choices"][0]["message"]["content"]) + del llm # free GPU memory before loading next + + +if __name__ == "__main__": + run("FALCON-MAMBA-7B-INSTRUCT (SSM)", MAMBA) + run("GEMMA-4-E4B-IT (Transformer)", GEMMA) diff --git a/experiments/mood_stream_poc.py b/experiments/mood_stream_poc.py new file mode 100644 index 0000000..502807f --- /dev/null +++ b/experiments/mood_stream_poc.py @@ -0,0 +1,91 @@ +""" +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/falcon-mamba-7B-instruct-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() diff --git a/experiments/recall_poc.py b/experiments/recall_poc.py new file mode 100644 index 0000000..67da52a --- /dev/null +++ b/experiments/recall_poc.py @@ -0,0 +1,212 @@ +""" +recall_poc.py — proof of concept for the 'recall' product idea. + +A daemon would normally feed events into Mamba state continuously. For the +POC we batch-feed the last N days of git activity from a set of repos, save +the resulting state, then in a SEPARATE invocation query the saved state. + +If the query in mode 2 produces an answer that genuinely reflects the input, +the concept is viable. If it produces vague/wrong answers — rethink. + +Usage: + # Step 1: ingest git history into a memba state + python recall_poc.py build --model + + # Step 2 (separate process — proves persistence): + python recall_poc.py query --model + python recall_poc.py query --model --q "Which project saw the most activity?" +""" + +from __future__ import annotations + +import argparse +import os +import subprocess +import sys +import time +from datetime import datetime, timedelta +from pathlib import Path + +# ── Config ─────────────────────────────────────────────────────── +SCAN_ROOTS = [ + Path.home() / "Desktop/Coding", + Path.home() / "Desktop/Coding/AI", +] +DAYS_BACK = 30 +STATE_DIR = Path.home() / ".recall/states" +SESSION_ID = "recall_poc" +DEFAULT_Q = ( + "Summarise what I've been working on over the last month. " + "Group by project. Mention the main themes per project. " + "Be concise — one short paragraph per project." +) + + +# ── Git scraping ───────────────────────────────────────────────── + +def find_repos(roots: list[Path]) -> list[Path]: + repos = [] + seen = set() + for root in roots: + if not root.exists(): + continue + for entry in sorted(root.iterdir()): + if not entry.is_dir() or not (entry / ".git").exists(): + continue + real = entry.resolve() + if real in seen: + continue + seen.add(real) + repos.append(entry) + return repos + + +def git_log(repo: Path, since_days: int) -> str: + """Return human-readable commit log for last `since_days` days.""" + try: + out = subprocess.check_output( + [ + "git", "-C", str(repo), "log", + f"--since={since_days} days ago", + "--no-merges", + "--date=short", + "--pretty=format:%ad %s", + ], + text=True, stderr=subprocess.DEVNULL, + ) + except subprocess.CalledProcessError: + return "" + return out.strip() + + +def build_corpus(repos: list[Path], since_days: int) -> tuple[str, int]: + """Build a single human-readable summary block. Returns (text, n_commits).""" + blocks: list[str] = [] + total = 0 + for repo in repos: + log = git_log(repo, since_days) + if not log: + continue + n = log.count("\n") + 1 + total += n + blocks.append(f"=== {repo.name} ({n} commits) ===\n{log}") + return "\n\n".join(blocks), total + + +# ── memba interaction ─────────────────────────────────────────── + +def open_session(model_path: str, n_gpu_layers: int) -> "Session": # noqa: F821 + # Imported lazily so query mode doesn't pay the cost when only listing + sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "python")) + from memba import Session + return Session( + model_path=model_path, + session_id=SESSION_ID, + state_dir=str(STATE_DIR), + n_gpu_layers=n_gpu_layers, + n_ctx=8192, # bigger than default so a month of commits fits comfortably + chat_format="chatml", + ) + + +# ── build mode ────────────────────────────────────────────────── + +def cmd_build(args: argparse.Namespace) -> int: + repos = find_repos(SCAN_ROOTS) + print(f"[build] discovered {len(repos)} git repos under {[str(r) for r in SCAN_ROOTS]}") + + corpus, n_commits = build_corpus(repos, args.days) + if not corpus: + print(f"[build] no commits in the last {args.days} days — abort", file=sys.stderr) + return 1 + + char_count = len(corpus) + print(f"[build] corpus: {n_commits} commits, {char_count:,} characters") + print(f"[build] first 400 chars:\n---\n{corpus[:400]}\n---") + + # Wipe any previous state for a clean test + state_file = STATE_DIR / f"{SESSION_ID}.memb" + if state_file.exists(): + state_file.unlink() + print(f"[build] removed previous state file") + + print(f"[build] loading model ({Path(args.model).name}) on GPU layers={args.gpu_layers}…") + t0 = time.time() + sess = open_session(args.model, args.gpu_layers) + print(f"[build] model loaded in {time.time()-t0:.1f}s, initial state {sess.state_size:,} B") + + prompt = ( + "I am going to give you my git commit history from the last month, " + "across several of my personal projects. Please READ it and remember " + "the overall picture — which projects I worked on, the kinds of " + "changes I made, and any themes that emerge. Just reply 'noted' when " + "you have processed it; I'll ask questions in a follow-up.\n\n" + f"=== git log (last {args.days} days, {n_commits} commits) ===\n\n" + f"{corpus}" + ) + + print(f"[build] feeding {len(prompt):,} characters into state…") + t0 = time.time() + ack = sess.chat(prompt, max_tokens=32) + elapsed = time.time() - t0 + print(f"[build] processed in {elapsed:.1f}s (~{len(prompt)/elapsed:,.0f} char/s)") + print(f"[build] model ack: {ack!r}") + print(f"[build] state after ingest: {sess.state_size:,} B") + + path = sess.save() + on_disk = path.stat().st_size + print(f"[build] state saved → {path} ({on_disk:,} B on disk)") + return 0 + + +# ── query mode ────────────────────────────────────────────────── + +def cmd_query(args: argparse.Namespace) -> int: + state_file = STATE_DIR / f"{SESSION_ID}.memb" + if not state_file.exists(): + print(f"[query] no state file at {state_file} — run `build` first", file=sys.stderr) + return 1 + + on_disk = state_file.stat().st_size + print(f"[query] state file: {state_file} ({on_disk:,} B)") + print(f"[query] loading model ({Path(args.model).name}) on GPU layers={args.gpu_layers}…") + t0 = time.time() + sess = open_session(args.model, args.gpu_layers) + print(f"[query] model+state loaded in {time.time()-t0:.1f}s, state size {sess.state_size:,} B") + + question = args.q or DEFAULT_Q + print(f"\n[query] question:\n{question}\n") + print(f"[query] response:\n---") + t0 = time.time() + reply = sess.chat(question, max_tokens=args.max_tokens) + print(reply) + print(f"--- ({time.time()-t0:.1f}s)") + return 0 + + +# ── main ──────────────────────────────────────────────────────── + +def main() -> int: + p = argparse.ArgumentParser(description="recall POC — git log → mamba state → cross-process query") + sub = p.add_subparsers(dest="cmd", required=True) + + common = argparse.ArgumentParser(add_help=False) + common.add_argument("--model", required=True, help="Path to GGUF model") + common.add_argument("--gpu-layers", type=int, default=-1, dest="gpu_layers", + help="Layers to offload (-1=all, 0=CPU)") + + b = sub.add_parser("build", parents=[common], help="Ingest git history into state") + b.add_argument("--days", type=int, default=DAYS_BACK) + b.set_defaults(func=cmd_build) + + q = sub.add_parser("query", parents=[common], help="Query the saved state") + q.add_argument("--q", help="Question to ask (default: monthly summary)") + q.add_argument("--max-tokens", type=int, default=512, dest="max_tokens") + q.set_defaults(func=cmd_query) + + args = p.parse_args() + return args.func(args) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/include/memba/state.h b/include/memba/state.h new file mode 100644 index 0000000..b6d9042 --- /dev/null +++ b/include/memba/state.h @@ -0,0 +1,87 @@ +#pragma once + +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* ── File format constants ─────────────────────────────────────── + * + * On-disk layout (little-endian throughout): + * + * [ magic : char[4] ] "MEMB" + * [ version : uint32 ] MEMBA_FILE_VERSION + * [ model_id : char[64] ] hex-encoded SHA-256 of first 1 KiB of GGUF + * [ n_ctx : uint32 ] llama_n_ctx() at save time + * [ llama_ver: uint32 ] reserved (0 for now) + * [ data_size: uint64 ] byte length of the opaque state blob + * [ data : uint8[] ] llama_state_get_data() blob + * [ crc32 : uint32 ] CRC-32 of data[] only + */ +#define MEMBA_FILE_MAGIC "MEMB" +#define MEMBA_FILE_VERSION 1u +#define MEMBA_MODEL_ID_LEN 64 + +/* ── Error codes ─────────────────────────────────────────────────*/ +#define MEMBA_OK 0 +#define MEMBA_ERR_IO -1 /* file open / read / write failed */ +#define MEMBA_ERR_MAGIC -2 /* bad magic bytes */ +#define MEMBA_ERR_VERSION -3 /* unsupported file version */ +#define MEMBA_ERR_MODEL_ID -4 /* model identity mismatch */ +#define MEMBA_ERR_CRC -5 /* CRC-32 checksum mismatch */ +#define MEMBA_ERR_ALLOC -6 /* memory allocation failed */ +#define MEMBA_ERR_CTX -7 /* null or invalid llama_context */ + +/* Opaque handle — one per llama_context you want to checkpoint. */ +typedef struct memba_state memba_state_t; + +struct llama_context; /* forward declaration */ + +/** + * Create a handle that wraps @p ctx. + * + * @param ctx Active llama_context. Must remain alive for the handle's + * entire lifetime. + * @param model_path Path to the GGUF file. Used to compute the model identity + * fingerprint embedded in every state file. Pass NULL to + * skip identity checking (strongly discouraged). + * @return New handle, or NULL on allocation failure. + */ +memba_state_t* memba_state_new(struct llama_context* ctx, const char* model_path); + +/** Free a handle created by memba_state_new(). Does NOT free ctx. */ +void memba_state_free(memba_state_t* state); + +/** + * Serialise the current SSM hidden state to @p path. + * + * Thread-safety: the caller must ensure no concurrent llama_decode() calls + * on @p ctx while this function executes. + * + * @return MEMBA_OK on success, negative error code on failure. + */ +int memba_state_save(memba_state_t* state, const char* path); + +/** + * Restore state from @p path into the wrapped llama_context. + * + * Validates magic, version, model_id, and CRC-32 before applying. + * + * @return MEMBA_OK on success, negative error code on failure. + */ +int memba_state_load(memba_state_t* state, const char* path); + +/** + * Return the serialised byte size of the current state (useful for logging). + * Returns 0 if the handle is NULL or ctx is invalid. + */ +size_t memba_state_get_size(memba_state_t* state); + +/** Human-readable description of an error code. Never returns NULL. */ +const char* memba_error_string(int err); + +#ifdef __cplusplus +} +#endif diff --git a/llama.cpp b/llama.cpp new file mode 160000 index 0000000..59778f0 --- /dev/null +++ b/llama.cpp @@ -0,0 +1 @@ +Subproject commit 59778f0196a82db32580bb649d5d839355d6d7bf diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..01e7d28 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,64 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "memba" +version = "0.1.0" +description = "Persistent memory layer for SSM-based LLMs (Falcon-Mamba, Zamba)" +readme = "README.md" +license = { text = "MIT" } +requires-python = ">=3.10" + +keywords = ["llm", "ssm", "mamba", "falcon-mamba", "memory", "llama-cpp"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", + "Topic :: Scientific/Engineering :: Artificial Intelligence", +] + +dependencies = [ + "llama-cpp-python>=0.2.0", + "typer>=0.9.0", + "rich>=13.0.0", + "pydantic>=2.0.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0", + "pytest-cov", + "ruff", + "mypy", +] + +[project.scripts] +memba = "memba.cli:main" + +[project.urls] +Homepage = "https://github.com/your-org/memba" +Repository = "https://github.com/your-org/memba" + +# ── Package discovery ──────────────────────────────────────────── +[tool.setuptools.packages.find] +where = ["python"] + +# ── Ruff (linting) ─────────────────────────────────────────────── +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "W", "I"] +ignore = ["E501"] + +# ── Mypy ──────────────────────────────────────────────────────── +[tool.mypy] +python_version = "3.10" +ignore_missing_imports = true diff --git a/python/memba/__init__.py b/python/memba/__init__.py new file mode 100644 index 0000000..53e6690 --- /dev/null +++ b/python/memba/__init__.py @@ -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" diff --git a/python/memba/cli.py b/python/memba/cli.py new file mode 100644 index 0000000..5df4ae5 --- /dev/null +++ b/python/memba/cli.py @@ -0,0 +1,257 @@ +""" +memba CLI — typer-based entry point. + +Sub-commands +------------ + memba chat --model --session [--gpu-layers N] → REPL + memba run --model --prompt [--save-state ] → one-shot + memba list [--state-dir ] → list sessions + memba rm [--state-dir ] → delete session + memba info [--state-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() diff --git a/python/memba/core.py b/python/memba/core.py new file mode 100644 index 0000000..5208065 --- /dev/null +++ b/python/memba/core.py @@ -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(" 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("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})" + ) diff --git a/src/cli.cpp b/src/cli.cpp new file mode 100644 index 0000000..fae1b4f --- /dev/null +++ b/src/cli.cpp @@ -0,0 +1,197 @@ +/* memba-cli — minimal SSM inference demo with state save/load + * + * Usage: + * memba-cli --model --prompt [--save-state ] + * memba-cli --model --load-state --prompt + * memba-cli --model --n-gpu-layers 35 --prompt + */ + +#include "memba/state.h" + +#include + +#include +#include +#include +#include +#include +#include + +/* ── Argument parsing ───────────────────────────────────────────*/ + +struct Args { + std::string model; + std::string prompt; + std::string save_state; + std::string load_state; + int n_predict = 256; + int n_gpu_layers = 0; + int n_ctx = 4096; + bool show_help = false; +}; + +static void print_usage(const char* prog) { + fprintf(stderr, + "Usage: %s --model --prompt [options]\n" + "\n" + "Options:\n" + " --model Path to GGUF model (required)\n" + " --prompt Input prompt (required unless --load-state)\n" + " --save-state Save SSM state after generation\n" + " --load-state Load SSM state before generation\n" + " --n-predict Tokens to generate (default: 256)\n" + " --n-gpu-layers Layers to offload to GPU (default: 0 = CPU)\n" + " --n-ctx Context size (default: 4096)\n" + " --help Show this message\n", + prog); +} + +static Args parse_args(int argc, char** argv) { + Args a; + for (int i = 1; i < argc; i++) { + auto next = [&]() -> const char* { + if (i + 1 >= argc) { + fprintf(stderr, "Missing value for %s\n", argv[i]); + exit(1); + } + return argv[++i]; + }; + std::string arg = argv[i]; + if (arg == "--model") a.model = next(); + else if (arg == "--prompt") a.prompt = next(); + else if (arg == "--save-state") a.save_state = next(); + else if (arg == "--load-state") a.load_state = next(); + else if (arg == "--n-predict") a.n_predict = std::atoi(next()); + else if (arg == "--n-gpu-layers") a.n_gpu_layers = std::atoi(next()); + else if (arg == "--n-ctx") a.n_ctx = std::atoi(next()); + else if (arg == "--help" || arg == "-h") a.show_help = true; + else { fprintf(stderr, "Unknown option: %s\n", argv[i]); exit(1); } + } + return a; +} + +/* ── Helpers ────────────────────────────────────────────────────*/ + +static void die(const char* msg) { + fprintf(stderr, "error: %s\n", msg); + exit(1); +} + +/* ── Main ───────────────────────────────────────────────────────*/ + +int main(int argc, char** argv) { + Args args = parse_args(argc, argv); + + if (args.show_help) { print_usage(argv[0]); return 0; } + if (args.model.empty()) die("--model is required"); + if (args.prompt.empty() && args.load_state.empty()) + die("--prompt or --load-state is required"); + + /* ── Initialise llama backend ───────────────────────────── */ + llama_backend_init(); + + /* ── Load model ─────────────────────────────────────────── */ + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = args.n_gpu_layers; + + llama_model* model = llama_model_load_from_file(args.model.c_str(), mparams); + if (!model) die("failed to load model"); + + /* ── Create context ─────────────────────────────────────── */ + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = (uint32_t)args.n_ctx; + + llama_context* ctx = llama_init_from_model(model, cparams); + if (!ctx) { llama_model_free(model); die("failed to create context"); } + + const llama_vocab* vocab = llama_model_get_vocab(model); + + /* ── memba handle ───────────────────────────────────────── */ + memba_state_t* mstate = memba_state_new(ctx, args.model.c_str()); + if (!mstate) { llama_free(ctx); llama_model_free(model); die("memba_state_new failed"); } + + /* ── Optionally restore a prior state ───────────────────── */ + if (!args.load_state.empty()) { + int rc = memba_state_load(mstate, args.load_state.c_str()); + if (rc != MEMBA_OK) { + fprintf(stderr, "load-state failed: %s\n", memba_error_string(rc)); + memba_state_free(mstate); llama_free(ctx); llama_model_free(model); + return 1; + } + fprintf(stderr, "[memba] state loaded from %s\n", args.load_state.c_str()); + } + + /* ── Tokenise prompt ────────────────────────────────────── */ + if (!args.prompt.empty()) { + const std::string& p = args.prompt; + + /* Over-allocate; llama_tokenize returns actual count */ + std::vector tokens(p.size() + 64); + int n_tokens = llama_tokenize( + vocab, p.c_str(), (int32_t)p.size(), + tokens.data(), (int32_t)tokens.size(), + /*add_special=*/true, /*parse_special=*/false); + + if (n_tokens < 0) { + /* Buffer was too small — resize and retry */ + tokens.resize((size_t)(-n_tokens)); + n_tokens = llama_tokenize( + vocab, p.c_str(), (int32_t)p.size(), + tokens.data(), (int32_t)tokens.size(), true, false); + } + if (n_tokens <= 0) die("tokenisation failed"); + tokens.resize((size_t)n_tokens); + + /* Prefill */ + llama_batch batch = llama_batch_get_one(tokens.data(), n_tokens); + if (llama_decode(ctx, batch) != 0) die("llama_decode (prefill) failed"); + } + + /* ── Decode / generate ──────────────────────────────────── */ + + /* Build a minimal greedy sampler chain */ + /* VERIFY: llama_sampler_chain_init / llama_sampler_chain_add may not exist + * in older builds — check llama.cpp/include/llama.h for sampler API. */ + llama_sampler* sampler = llama_sampler_chain_init(llama_sampler_chain_default_params()); + llama_sampler_chain_add(sampler, llama_sampler_init_greedy()); + + fprintf(stderr, "\n"); + for (int i = 0; i < args.n_predict; i++) { + llama_token tok = llama_sampler_sample(sampler, ctx, -1); + if (llama_vocab_is_eog(vocab, tok)) break; + + char piece[256]; + int n = llama_token_to_piece(vocab, tok, piece, sizeof(piece), 0, false); + if (n > 0) { + fwrite(piece, 1, (size_t)n, stdout); + fflush(stdout); + } + + llama_sampler_accept(sampler, tok); + + llama_batch next = llama_batch_get_one(&tok, 1); + if (llama_decode(ctx, next) != 0) break; + } + fprintf(stdout, "\n"); + + llama_sampler_free(sampler); + + /* ── Optionally persist state ───────────────────────────── */ + if (!args.save_state.empty()) { + int rc = memba_state_save(mstate, args.save_state.c_str()); + if (rc != MEMBA_OK) { + fprintf(stderr, "save-state failed: %s\n", memba_error_string(rc)); + } else { + fprintf(stderr, "[memba] state saved to %s (%zu bytes)\n", + args.save_state.c_str(), memba_state_get_size(mstate)); + } + } + + /* ── Cleanup ────────────────────────────────────────────── */ + memba_state_free(mstate); + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + + return 0; +} diff --git a/src/state.cpp b/src/state.cpp new file mode 100644 index 0000000..f3357ed --- /dev/null +++ b/src/state.cpp @@ -0,0 +1,299 @@ +#include "memba/state.h" + +#include + +#include +#include +#include +#include +#include +#include + +/* ═══════════════════════════════════════════════════════════════ + * CRC-32 (ISO 3309 polynomial 0xEDB88320, no table, branch-free) + * ═══════════════════════════════════════════════════════════════ */ + +static uint32_t crc32_compute(const uint8_t* data, size_t len) { + uint32_t crc = 0xFFFFFFFFu; + for (size_t i = 0; i < len; i++) { + crc ^= data[i]; + for (int j = 0; j < 8; j++) { + uint32_t mask = static_cast(-(int32_t)(crc & 1u)); + crc = (crc >> 1u) ^ (0xEDB88320u & mask); + } + } + return ~crc; +} + +/* ═══════════════════════════════════════════════════════════════ + * Minimal SHA-256 (public domain, derived from Brad Conte's work) + * ═══════════════════════════════════════════════════════════════ */ + +#define ROTR32(x, n) (((x) >> (n)) | ((x) << (32u - (n)))) +#define CH(x, y, z) (((x) & (y)) ^ (~(x) & (z))) +#define MAJ(x, y, z) (((x) & (y)) ^ ((x) & (z)) ^ ((y) & (z))) +#define EP0(x) (ROTR32(x, 2) ^ ROTR32(x, 13) ^ ROTR32(x, 22)) +#define EP1(x) (ROTR32(x, 6) ^ ROTR32(x, 11) ^ ROTR32(x, 25)) +#define SIG0(x) (ROTR32(x, 7) ^ ROTR32(x, 18) ^ ((x) >> 3u)) +#define SIG1(x) (ROTR32(x, 17) ^ ROTR32(x, 19) ^ ((x) >> 10u)) + +static const uint32_t K256[64] = { + 0x428a2f98,0x71374491,0xb5c0fbcf,0xe9b5dba5,0x3956c25b,0x59f111f1,0x923f82a4,0xab1c5ed5, + 0xd807aa98,0x12835b01,0x243185be,0x550c7dc3,0x72be5d74,0x80deb1fe,0x9bdc06a7,0xc19bf174, + 0xe49b69c1,0xefbe4786,0x0fc19dc6,0x240ca1cc,0x2de92c6f,0x4a7484aa,0x5cb0a9dc,0x76f988da, + 0x983e5152,0xa831c66d,0xb00327c8,0xbf597fc7,0xc6e00bf3,0xd5a79147,0x06ca6351,0x14292967, + 0x27b70a85,0x2e1b2138,0x4d2c6dfc,0x53380d13,0x650a7354,0x766a0abb,0x81c2c92e,0x92722c85, + 0xa2bfe8a1,0xa81a664b,0xc24b8b70,0xc76c51a3,0xd192e819,0xd6990624,0xf40e3585,0x106aa070, + 0x19a4c116,0x1e376c08,0x2748774c,0x34b0bcb5,0x391c0cb3,0x4ed8aa4a,0x5b9cca4f,0x682e6ff3, + 0x748f82ee,0x78a5636f,0x84c87814,0x8cc70208,0x90befffa,0xa4506ceb,0xbef9a3f7,0xc67178f2, +}; + +struct SHA256Ctx { + uint8_t buf[64]; + uint32_t buflen; + uint64_t bitlen; + uint32_t state[8]; +}; + +static void sha256_transform(SHA256Ctx* c, const uint8_t* d) { + uint32_t m[64], a, b, cc, dd, e, f, g, h, t1, t2; + for (uint32_t i = 0, j = 0; i < 16; i++, j += 4) + m[i] = ((uint32_t)d[j]<<24)|((uint32_t)d[j+1]<<16)|((uint32_t)d[j+2]<<8)|(uint32_t)d[j+3]; + for (uint32_t i = 16; i < 64; i++) + m[i] = SIG1(m[i-2]) + m[i-7] + SIG0(m[i-15]) + m[i-16]; + a=c->state[0]; b=c->state[1]; cc=c->state[2]; dd=c->state[3]; + e=c->state[4]; f=c->state[5]; g=c->state[6]; h=c->state[7]; + for (uint32_t i = 0; i < 64; i++) { + t1 = h + EP1(e) + CH(e,f,g) + K256[i] + m[i]; + t2 = EP0(a) + MAJ(a,b,cc); + h=g; g=f; f=e; e=dd+t1; dd=cc; cc=b; b=a; a=t1+t2; + } + c->state[0]+=a; c->state[1]+=b; c->state[2]+=cc; c->state[3]+=dd; + c->state[4]+=e; c->state[5]+=f; c->state[6]+=g; c->state[7]+=h; +} + +static void sha256_init(SHA256Ctx* c) { + c->buflen = 0; c->bitlen = 0; + c->state[0]=0x6a09e667; c->state[1]=0xbb67ae85; + c->state[2]=0x3c6ef372; c->state[3]=0xa54ff53a; + c->state[4]=0x510e527f; c->state[5]=0x9b05688c; + c->state[6]=0x1f83d9ab; c->state[7]=0x5be0cd19; +} + +static void sha256_update(SHA256Ctx* c, const uint8_t* data, size_t len) { + for (size_t i = 0; i < len; i++) { + c->buf[c->buflen++] = data[i]; + if (c->buflen == 64) { + sha256_transform(c, c->buf); + c->bitlen += 512; + c->buflen = 0; + } + } +} + +static void sha256_final(SHA256Ctx* c, uint8_t hash[32]) { + uint32_t i = c->buflen; + c->buf[i++] = 0x80; + if (c->buflen < 56) { + while (i < 56) c->buf[i++] = 0; + } else { + while (i < 64) c->buf[i++] = 0; + sha256_transform(c, c->buf); + memset(c->buf, 0, 56); + } + c->bitlen += (uint64_t)c->buflen * 8; + for (int k = 7; k >= 0; k--) + c->buf[56 + (7-k)] = (uint8_t)(c->bitlen >> (k * 8)); + sha256_transform(c, c->buf); + for (i = 0; i < 4; i++) { + hash[i] = (c->state[0] >> (24 - i*8)) & 0xFF; + hash[i+4] = (c->state[1] >> (24 - i*8)) & 0xFF; + hash[i+8] = (c->state[2] >> (24 - i*8)) & 0xFF; + hash[i+12] = (c->state[3] >> (24 - i*8)) & 0xFF; + hash[i+16] = (c->state[4] >> (24 - i*8)) & 0xFF; + hash[i+20] = (c->state[5] >> (24 - i*8)) & 0xFF; + hash[i+24] = (c->state[6] >> (24 - i*8)) & 0xFF; + hash[i+28] = (c->state[7] >> (24 - i*8)) & 0xFF; + } +} + +/* ═══════════════════════════════════════════════════════════════ + * Model-ID computation + * ═══════════════════════════════════════════════════════════════ */ + +static void compute_model_id(const char* model_path, char out[MEMBA_MODEL_ID_LEN]) { + memset(out, 0, MEMBA_MODEL_ID_LEN); + + uint8_t raw[1024] = {}; + size_t nread = 0; + + if (model_path) { + FILE* fp = fopen(model_path, "rb"); + if (fp) { + nread = fread(raw, 1, sizeof(raw), fp); + fclose(fp); + } + } + if (nread == 0 && model_path) { + /* Fallback: hash the path string itself */ + const char* p = model_path; + while (*p) raw[nread++ % sizeof(raw)] ^= (uint8_t)*p++; + if (nread == 0) nread = 1; + } + + SHA256Ctx ctx; + sha256_init(&ctx); + sha256_update(&ctx, raw, nread); + uint8_t digest[32]; + sha256_final(&ctx, digest); + + static const char hex[] = "0123456789abcdef"; + for (int i = 0; i < 32; i++) { + out[i*2] = hex[digest[i] >> 4]; + out[i*2+1] = hex[digest[i] & 0x0F]; + } + /* 32*2 == 64 == MEMBA_MODEL_ID_LEN, no truncation needed */ +} + +/* ═══════════════════════════════════════════════════════════════ + * llama.cpp state API (post-2024 names) + * ═══════════════════════════════════════════════════════════════ */ + +static size_t state_get_size(struct llama_context* ctx) { + return llama_state_get_size(ctx); +} + +static size_t state_get_data(struct llama_context* ctx, uint8_t* dst, size_t size) { + return llama_state_get_data(ctx, dst, size); +} + +static size_t state_set_data(struct llama_context* ctx, const uint8_t* src, size_t size) { + return llama_state_set_data(ctx, src, size); +} + +/* ═══════════════════════════════════════════════════════════════ + * Public struct + * ═══════════════════════════════════════════════════════════════ */ + +struct memba_state { + struct llama_context* ctx; + char model_id[MEMBA_MODEL_ID_LEN]; +}; + +/* ═══════════════════════════════════════════════════════════════ + * Public API + * ═══════════════════════════════════════════════════════════════ */ + +memba_state_t* memba_state_new(struct llama_context* ctx, const char* model_path) { + if (!ctx) return nullptr; + auto* s = static_cast(malloc(sizeof(memba_state_t))); + if (!s) return nullptr; + s->ctx = ctx; + compute_model_id(model_path, s->model_id); + return s; +} + +void memba_state_free(memba_state_t* state) { + free(state); +} + +size_t memba_state_get_size(memba_state_t* state) { + if (!state || !state->ctx) return 0; + return state_get_size(state->ctx); +} + +int memba_state_save(memba_state_t* state, const char* path) { + if (!state || !state->ctx) return MEMBA_ERR_CTX; + if (!path) return MEMBA_ERR_IO; + + size_t sz = state_get_size(state->ctx); + if (sz == 0) return MEMBA_ERR_CTX; + + std::vector blob(sz); + size_t written = state_get_data(state->ctx, blob.data(), sz); + if (written == 0) return MEMBA_ERR_CTX; + blob.resize(written); + + uint32_t crc = crc32_compute(blob.data(), blob.size()); + uint32_t version = MEMBA_FILE_VERSION; + uint32_t n_ctx = (uint32_t)llama_n_ctx(state->ctx); + uint32_t llama_ver = 0; /* reserved */ + uint64_t data_size = (uint64_t)blob.size(); + + FILE* fp = fopen(path, "wb"); + if (!fp) return MEMBA_ERR_IO; + + bool ok = true; + ok = ok && fwrite(MEMBA_FILE_MAGIC, 1, 4, fp) == 4; + ok = ok && fwrite(&version, 4, 1, fp) == 1; + ok = ok && fwrite(state->model_id, 1, MEMBA_MODEL_ID_LEN, fp) == (size_t)MEMBA_MODEL_ID_LEN; + ok = ok && fwrite(&n_ctx, 4, 1, fp) == 1; + ok = ok && fwrite(&llama_ver, 4, 1, fp) == 1; + ok = ok && fwrite(&data_size, 8, 1, fp) == 1; + ok = ok && fwrite(blob.data(), 1, blob.size(), fp) == blob.size(); + ok = ok && fwrite(&crc, 4, 1, fp) == 1; + + fclose(fp); + return ok ? MEMBA_OK : MEMBA_ERR_IO; +} + +int memba_state_load(memba_state_t* state, const char* path) { + if (!state || !state->ctx) return MEMBA_ERR_CTX; + if (!path) return MEMBA_ERR_IO; + + FILE* fp = fopen(path, "rb"); + if (!fp) return MEMBA_ERR_IO; + + char magic[4]; + uint32_t version, n_ctx, llama_ver; + char file_model_id[MEMBA_MODEL_ID_LEN]; + uint64_t data_size; + + bool ok = true; + ok = ok && fread(magic, 1, 4, fp) == 4; + ok = ok && fread(&version, 4, 1, fp) == 1; + ok = ok && fread(file_model_id, 1, MEMBA_MODEL_ID_LEN, fp) == (size_t)MEMBA_MODEL_ID_LEN; + ok = ok && fread(&n_ctx, 4, 1, fp) == 1; + ok = ok && fread(&llama_ver, 4, 1, fp) == 1; + ok = ok && fread(&data_size, 8, 1, fp) == 1; + + if (!ok) { fclose(fp); return MEMBA_ERR_IO; } + + if (memcmp(magic, MEMBA_FILE_MAGIC, 4) != 0) { fclose(fp); return MEMBA_ERR_MAGIC; } + if (version != MEMBA_FILE_VERSION) { fclose(fp); return MEMBA_ERR_VERSION; } + if (memcmp(file_model_id, state->model_id, MEMBA_MODEL_ID_LEN) != 0) { + fclose(fp); return MEMBA_ERR_MODEL_ID; + } + + if (data_size > (uint64_t)SIZE_MAX) { fclose(fp); return MEMBA_ERR_ALLOC; } + std::vector blob((size_t)data_size); + + if (fread(blob.data(), 1, blob.size(), fp) != blob.size()) { + fclose(fp); return MEMBA_ERR_IO; + } + + uint32_t stored_crc; + if (fread(&stored_crc, 4, 1, fp) != 1) { fclose(fp); return MEMBA_ERR_IO; } + fclose(fp); + + if (crc32_compute(blob.data(), blob.size()) != stored_crc) return MEMBA_ERR_CRC; + + size_t restored = state_set_data(state->ctx, blob.data(), blob.size()); + if (restored == 0) return MEMBA_ERR_CTX; + + return MEMBA_OK; +} + +const char* memba_error_string(int err) { + switch (err) { + case MEMBA_OK: return "OK"; + case MEMBA_ERR_IO: return "I/O error"; + case MEMBA_ERR_MAGIC: return "bad file magic (not a memba state file)"; + case MEMBA_ERR_VERSION: return "unsupported file version"; + case MEMBA_ERR_MODEL_ID: return "model identity mismatch"; + case MEMBA_ERR_CRC: return "CRC-32 checksum mismatch (file corrupted)"; + case MEMBA_ERR_ALLOC: return "memory allocation failed"; + case MEMBA_ERR_CTX: return "invalid or null llama_context"; + default: return "unknown error"; + } +}