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
+45
@@ -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
|
||||
@@ -0,0 +1,3 @@
|
||||
[submodule "llama.cpp"]
|
||||
path = llama.cpp
|
||||
url = https://github.com/ggerganov/llama.cpp.git
|
||||
@@ -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)
|
||||
@@ -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 <memba/state.h>
|
||||
#include <llama.h>
|
||||
|
||||
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 <session> <new-name>` — 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.
|
||||
@@ -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 ."
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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.
|
||||
@@ -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]()
|
||||
@@ -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]()
|
||||
@@ -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)
|
||||
@@ -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()
|
||||
@@ -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 <gguf>
|
||||
|
||||
# Step 2 (separate process — proves persistence):
|
||||
python recall_poc.py query --model <gguf>
|
||||
python recall_poc.py query --model <gguf> --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())
|
||||
@@ -0,0 +1,87 @@
|
||||
#pragma once
|
||||
|
||||
#include <stddef.h>
|
||||
#include <stdint.h>
|
||||
|
||||
#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
|
||||
Submodule
+1
Submodule llama.cpp added at 59778f0196
@@ -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
|
||||
@@ -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})"
|
||||
)
|
||||
+197
@@ -0,0 +1,197 @@
|
||||
/* memba-cli — minimal SSM inference demo with state save/load
|
||||
*
|
||||
* Usage:
|
||||
* memba-cli --model <gguf> --prompt <text> [--save-state <file>]
|
||||
* memba-cli --model <gguf> --load-state <file> --prompt <text>
|
||||
* memba-cli --model <gguf> --n-gpu-layers 35 --prompt <text>
|
||||
*/
|
||||
|
||||
#include "memba/state.h"
|
||||
|
||||
#include <llama.h>
|
||||
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
|
||||
/* ── 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 <gguf> --prompt <text> [options]\n"
|
||||
"\n"
|
||||
"Options:\n"
|
||||
" --model <path> Path to GGUF model (required)\n"
|
||||
" --prompt <text> Input prompt (required unless --load-state)\n"
|
||||
" --save-state <path> Save SSM state after generation\n"
|
||||
" --load-state <path> Load SSM state before generation\n"
|
||||
" --n-predict <N> Tokens to generate (default: 256)\n"
|
||||
" --n-gpu-layers<N> Layers to offload to GPU (default: 0 = CPU)\n"
|
||||
" --n-ctx <N> 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<llama_token> 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;
|
||||
}
|
||||
+299
@@ -0,0 +1,299 @@
|
||||
#include "memba/state.h"
|
||||
|
||||
#include <llama.h>
|
||||
|
||||
#include <cassert>
|
||||
#include <cstdio>
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <cstdint>
|
||||
#include <vector>
|
||||
|
||||
/* ═══════════════════════════════════════════════════════════════
|
||||
* 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<uint32_t>(-(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<memba_state_t*>(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<uint8_t> 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<uint8_t> 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";
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user