pylolm 0.1.0__py3-none-any.whl

pylolm/__init__.py

@@ -0,0 +1,36 @@
+"""pylolm (LoLm): run GGUF language models on-device on Android, from Python.
+
+An Android-friendly take on llama-cpp-python, built on llama.cpp and designed for
+Kivy, Flet and Termux.
+"""
+
+from .errors import (
+    ContextOverflowError,
+    DownloadError,
+    GenerationError,
+    LibraryNotFoundError,
+    ModelLoadError,
+    PyLoLmError,
+)
+from .llama import Llama
+
+try:
+    from importlib.metadata import PackageNotFoundError, version as _version
+
+    try:
+        __version__ = _version("pylolm")
+    except PackageNotFoundError:  # running from a source checkout
+        __version__ = "0.0.0+local"
+except ImportError:  # pragma: no cover
+    __version__ = "0.0.0+local"
+
+__all__ = [
+    "Llama",
+    "PyLoLmError",
+    "LibraryNotFoundError",
+    "ModelLoadError",
+    "ContextOverflowError",
+    "DownloadError",
+    "GenerationError",
+    "__version__",
+]

pylolm/__main__.py

@@ -0,0 +1,168 @@
+"""Command line: ``python -m pylolm {doctor,get-libs,build-native,chat}`` (also ``pylolm``)."""
+
+import argparse
+import platform
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+from . import __version__, _lib
+from .errors import DownloadError, PyLoLmError
+
+
+def cmd_doctor(_args) -> int:
+    print("pylolm %s" % __version__)
+    print("python %s on %s (%s)" % (platform.python_version(), platform.system(), platform.machine()))
+    print("search order for native libraries:")
+    for d in _lib.candidate_dirs():
+        found = (d / _lib.SHIM_NAME).is_file()
+        print("  [%s] %s" % ("x" if found else " ", d))
+    try:
+        _lib.load()
+    except PyLoLmError as exc:
+        print("\nNOT READY: %s" % exc)
+        return 1
+    print("\nOK: native libraries loaded from %s" % _lib.loaded_dir())
+    return 0
+
+
+def cmd_get_libs(args) -> int:
+    from . import _download
+
+    pending = [False]  # a progress line is on screen without a trailing newline
+
+    def progress(done, total):
+        pending[0] = True
+        mb = done / 1e6
+        text = "  %.1f / %.1f MB" % (mb, total / 1e6) if total else "  %.1f MB" % mb
+        print("\r" + text, end="", flush=True)
+
+    def log(message):
+        if pending[0]:
+            print()
+            pending[0] = False
+        print(message)
+
+    try:
+        _download.get_libs(
+            version=args.version,
+            repo=args.repo,
+            base_url=args.base_url,
+            force=args.force,
+            progress=progress,
+            log=log,
+        )
+    except DownloadError as exc:
+        if pending[0]:
+            print()
+        print("error: %s" % exc, file=sys.stderr)
+        return 1
+    print("\nNext: pylolm doctor")
+    return 0
+
+
+def cmd_build_native(_args) -> int:
+    script = Path(__file__).resolve().parent / "native" / "build_termux.sh"
+    if not script.is_file():
+        print("build script not found: %s" % script, file=sys.stderr)
+        return 1
+    return subprocess.call(["bash", str(script)])
+
+
+def cmd_chat(args) -> int:
+    from .llama import Llama
+
+    try:
+        llm = Llama(args.model, n_ctx=args.n_ctx, n_threads=args.threads)
+    except (PyLoLmError, FileNotFoundError) as exc:
+        print("error: %s" % exc, file=sys.stderr)
+        return 1
+
+    history = []
+    if args.system:
+        history.append({"role": "system", "content": args.system})
+    print("pylolm chat - /reset clears history, /exit quits.\n")
+
+    while True:
+        try:
+            user = input("you> ").strip()
+        except (EOFError, KeyboardInterrupt):
+            print()
+            break
+        if not user:
+            continue
+        if user == "/exit":
+            break
+        if user == "/reset":
+            history = history[:1] if args.system else []
+            llm.reset()
+            print("(history cleared)")
+            continue
+
+        history.append({"role": "user", "content": user})
+        reply = []
+        t0 = time.time()
+        n = 0
+        try:
+            print("bot> ", end="", flush=True)
+            for chunk in llm.create_chat_completion(
+                history, max_tokens=args.max_tokens, temperature=args.temp, stream=True
+            ):
+                text = chunk["choices"][0]["delta"].get("content")
+                if text:
+                    n += 1
+                    reply.append(text)
+                    print(text, end="", flush=True)
+        except KeyboardInterrupt:
+            llm.stop()
+            print("\n(interrupted)")
+        except PyLoLmError as exc:
+            print("\nerror: %s" % exc)
+            history.pop()
+            continue
+        dt = max(time.time() - t0, 1e-6)
+        print("\n[%d chunks, %.1f/s]\n" % (n, n / dt))
+        history.append({"role": "assistant", "content": "".join(reply)})
+
+    llm.close()
+    return 0
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(prog="pylolm", description="Local LLMs on Android, from Python")
+    parser.add_argument("--version", action="version", version="pylolm " + __version__)
+    sub = parser.add_subparsers(dest="command")
+
+    sub.add_parser("doctor", help="check that the native libraries can be found and loaded")
+    sub.add_parser("build-native", help="build llama.cpp + the shim on this device (slow, Termux)")
+
+    get = sub.add_parser("get-libs", help="download the prebuilt native libraries (fast)")
+    get.add_argument("--version", default=None, help="release to download (default: this pylolm)")
+    get.add_argument("--repo", default=None, help="GitHub repo as owner/name")
+    get.add_argument("--base-url", default=None, help="download from this URL instead of GitHub")
+    get.add_argument("--force", action="store_true", help="replace existing libraries")
+
+    chat = sub.add_parser("chat", help="interactive chat in the terminal")
+    chat.add_argument("model", help="path to a .gguf file")
+    chat.add_argument("--n-ctx", type=int, default=2048)
+    chat.add_argument("--threads", type=int, default=None)
+    chat.add_argument("--system", default=None, help="system prompt")
+    chat.add_argument("--max-tokens", type=int, default=512)
+    chat.add_argument("--temp", type=float, default=0.7)
+
+    args = parser.parse_args(argv)
+    handlers = {
+        "doctor": cmd_doctor,
+        "get-libs": cmd_get_libs,
+        "build-native": cmd_build_native,
+        "chat": cmd_chat,
+    }
+    if args.command not in handlers:
+        parser.print_help()
+        return 0
+    return handlers[args.command](args)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

pylolm/_download.py

@@ -0,0 +1,184 @@
+"""Download the prebuilt native libraries from a GitHub Release (``pylolm get-libs``).
+
+Safety: the archive is only used if its SHA-256 matches the checksum file published next to
+it, and extraction only accepts plain ``*.so`` files (no paths, links or other entry types).
+"""
+
+import hashlib
+import os
+import platform
+import shutil
+import tarfile
+import tempfile
+import urllib.error
+import urllib.request
+from pathlib import Path
+from typing import Callable, List, Optional
+
+from .errors import DownloadError
+
+ASSET = "pylolm-libs-android-arm64-v8a.tar.gz"
+DEFAULT_REPO = "auraway/pylolm"
+SHIM_NAME = "libpylolm_shim.so"
+MAX_BYTES = 300 * 1024 * 1024  # refuse absurdly large downloads
+
+Progress = Callable[[int, int], None]
+
+
+def default_dest() -> Path:
+    """Same place the loader and `pylolm build-native` use."""
+    env = os.environ.get("PYLOLM_LIB_DIR")
+    return Path(env).expanduser() if env else Path.home() / ".pylolm" / "lib"
+
+
+def release_base_url(version: str, repo: Optional[str] = None) -> str:
+    repo = repo or os.environ.get("PYLOLM_RELEASE_REPO") or DEFAULT_REPO
+    return "https://github.com/%s/releases/download/v%s" % (repo, version)
+
+
+def _fetch(url: str, target: Path, progress: Optional[Progress] = None) -> None:
+    req = urllib.request.Request(url, headers={"User-Agent": "pylolm-get-libs"})
+    try:
+        with urllib.request.urlopen(req, timeout=30) as resp, open(str(target), "wb") as out:
+            total = int(resp.headers.get("Content-Length") or 0)
+            done = 0
+            while True:
+                block = resp.read(1 << 16)
+                if not block:
+                    break
+                done += len(block)
+                if done > MAX_BYTES:
+                    raise DownloadError("download is larger than expected; aborting")
+                out.write(block)
+                if progress:
+                    progress(done, total)
+    except urllib.error.HTTPError as exc:
+        if exc.code == 404:
+            raise DownloadError(
+                "not found (404): %s\n"
+                "Has this version been released with prebuilt libraries yet? "
+                "Use `pylolm build-native` instead." % url
+            )
+        raise DownloadError("HTTP error %d for %s" % (exc.code, url))
+    except (urllib.error.URLError, OSError) as exc:
+        raise DownloadError("could not download %s: %s" % (url, exc))
+
+
+def _sha256_of(path: Path) -> str:
+    h = hashlib.sha256()
+    with open(str(path), "rb") as f:
+        for block in iter(lambda: f.read(1 << 20), b""):
+            h.update(block)
+    return h.hexdigest()
+
+
+def _expected_sha256(path: Path) -> str:
+    try:
+        token = path.read_text().split()[0].lower()
+    except (IndexError, UnicodeDecodeError, OSError):
+        raise DownloadError("checksum file is empty or unreadable")
+    if len(token) != 64 or any(c not in "0123456789abcdef" for c in token):
+        raise DownloadError("checksum file does not contain a SHA-256 hash")
+    return token
+
+
+def _safe_extract(archive: Path, out_dir: Path) -> List[str]:
+    """Extract plain *.so files only, writing them ourselves (no tar path handling)."""
+    names: List[str] = []
+    try:
+        with tarfile.open(str(archive), "r:gz") as tar:
+            members = tar.getmembers()
+            if not members:
+                raise DownloadError("archive is empty")
+            for m in members:
+                name = m.name
+                if (
+                    not m.isfile()
+                    or os.path.basename(name) != name
+                    or name.startswith(".")
+                    or not name.endswith(".so")
+                ):
+                    raise DownloadError("unexpected entry in archive: %r" % name)
+            for m in members:
+                src = tar.extractfile(m)
+                if src is None:
+                    raise DownloadError("could not read %r from archive" % m.name)
+                target = out_dir / m.name
+                with open(str(target), "wb") as dst:
+                    shutil.copyfileobj(src, dst)
+                os.chmod(str(target), 0o755)
+                names.append(m.name)
+    except (tarfile.TarError, EOFError, OSError) as exc:
+        raise DownloadError("archive is corrupt: %s" % exc)
+    if SHIM_NAME not in names:
+        raise DownloadError("archive does not contain %s" % SHIM_NAME)
+    return names
+
+
+def get_libs(
+    version: Optional[str] = None,
+    repo: Optional[str] = None,
+    base_url: Optional[str] = None,
+    dest: Optional[Path] = None,
+    force: bool = False,
+    progress: Optional[Progress] = None,
+    log: Callable[[str], None] = print,
+) -> Path:
+    """Download, verify and install the prebuilt libraries. Returns the install directory."""
+    machine = platform.machine().lower()
+    if machine not in ("aarch64", "arm64") and not force:
+        raise DownloadError(
+            "the prebuilt libraries are for Android arm64 (aarch64), but this device "
+            "reports %r. Use --force to install anyway, or run `pylolm build-native`." % machine
+        )
+
+    if base_url is None:
+        if version is None:
+            from . import __version__
+
+            if "+" in __version__:
+                raise DownloadError(
+                    "this looks like a source checkout (version %s); pass --version X.Y.Z"
+                    % __version__
+                )
+            version = __version__
+        base_url = release_base_url(version, repo)
+    base_url = base_url.rstrip("/")
+
+    dest = Path(dest) if dest else default_dest()
+    if (dest / SHIM_NAME).is_file() and not force:
+        raise DownloadError(
+            "native libraries are already installed in %s. Use --force to replace them." % dest
+        )
+
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    staging = Path(tempfile.mkdtemp(prefix=".pylolm-get-libs-", dir=str(dest.parent)))
+    try:
+        archive = staging / ASSET
+        checksum = staging / (ASSET + ".sha256")
+
+        log("Downloading %s ..." % ASSET)
+        _fetch("%s/%s" % (base_url, ASSET), archive, progress)
+        _fetch("%s/%s.sha256" % (base_url, ASSET), checksum)
+
+        log("Verifying checksum ...")
+        expected, actual = _expected_sha256(checksum), _sha256_of(archive)
+        if expected != actual:
+            raise DownloadError(
+                "checksum mismatch (expected %s..., got %s...); the file was NOT installed"
+                % (expected[:12], actual[:12])
+            )
+
+        unpacked = staging / "unpacked"
+        unpacked.mkdir()
+        names = _safe_extract(archive, unpacked)
+
+        dest.mkdir(parents=True, exist_ok=True)
+        for old in dest.glob("*.so"):
+            old.unlink()
+        for name in names:
+            shutil.move(str(unpacked / name), str(dest / name))
+        log("Installed %d libraries to %s" % (len(names), dest))
+        return dest
+    finally:
+        shutil.rmtree(str(staging), ignore_errors=True)

pylolm/_lib.py

@@ -0,0 +1,155 @@
+"""Locate and load the native libraries, and declare the ctypes signatures of the shim."""
+
+import ctypes
+import os
+import threading
+from pathlib import Path
+from typing import Iterator, List
+
+from .errors import LibraryNotFoundError
+
+SHIM_ABI_VERSION = 1
+SHIM_NAME = "libpylolm_shim.so"
+
+# Loaded in this order with RTLD_GLOBAL so the dynamic linker finds them by soname.
+# (name, required)
+_DEPENDENCIES = [
+    ("libggml-base.so", True),
+    ("libggml-cpu.so", False),  # only present when CPU variants are NOT used
+    ("libggml.so", True),
+    ("libllama.so", True),
+]
+
+# void callback(const char *piece, int len, void *user) -> int (non-zero = stop)
+TOKEN_CB = ctypes.CFUNCTYPE(ctypes.c_int, ctypes.c_void_p, ctypes.c_int, ctypes.c_void_p)
+
+_lock = threading.Lock()
+_shim = None
+_lib_dir = None
+
+
+def candidate_dirs() -> Iterator[Path]:
+    """Directories searched for the native libraries, in priority order."""
+    env = os.environ.get("PYLOLM_LIB_DIR")
+    if env:
+        yield Path(env).expanduser()
+    yield Path(__file__).resolve().parent / "lib"
+    try:
+        yield Path.home() / ".pylolm" / "lib"
+    except RuntimeError:  # no resolvable home directory
+        pass
+
+
+def find_lib_dir() -> Path:
+    searched: List[str] = []
+    for d in candidate_dirs():
+        searched.append(str(d))
+        if (d / SHIM_NAME).is_file():
+            return d
+    raise LibraryNotFoundError(
+        "pylolm native libraries not found (looked for %s in: %s).\n"
+        "  - Download the prebuilt libraries (fast): pylolm get-libs\n"
+        "  - Or build them on this device (slow, ~30 min): pylolm build-native\n"
+        "  - Or point PYLOLM_LIB_DIR at a folder containing the .so files."
+        % (SHIM_NAME, ", ".join(searched))
+    )
+
+
+def _bind(shim: ctypes.CDLL) -> None:
+    c = ctypes
+    sampling = [
+        c.c_float,  # temperature
+        c.c_float,  # top_p
+        c.c_int,  # top_k
+        c.c_float,  # min_p
+        c.c_float,  # repeat_penalty
+        c.c_uint32,  # seed
+    ]
+
+    shim.shim_abi_version.argtypes = []
+    shim.shim_abi_version.restype = c.c_int
+
+    shim.shim_set_verbose.argtypes = [c.c_int]
+    shim.shim_set_verbose.restype = None
+
+    shim.shim_init.argtypes = [c.c_char_p]
+    shim.shim_init.restype = c.c_int
+
+    shim.shim_load.argtypes = [c.c_char_p, c.c_int, c.c_int]
+    shim.shim_load.restype = c.c_void_p
+
+    shim.shim_stop.argtypes = [c.c_void_p]
+    shim.shim_stop.restype = None
+
+    shim.shim_reset.argtypes = [c.c_void_p]
+    shim.shim_reset.restype = None
+
+    shim.shim_n_ctx.argtypes = [c.c_void_p]
+    shim.shim_n_ctx.restype = c.c_int
+
+    shim.shim_count_tokens.argtypes = [c.c_void_p, c.c_char_p, c.c_int]
+    shim.shim_count_tokens.restype = c.c_int
+
+    shim.shim_last_stats.argtypes = [c.c_void_p, c.POINTER(c.c_int)]
+    shim.shim_last_stats.restype = None
+
+    shim.shim_chat.argtypes = (
+        [
+            c.c_void_p,  # handle
+            c.POINTER(c.c_char_p),  # roles
+            c.POINTER(c.c_char_p),  # contents
+            c.c_int,  # n_msgs
+            c.c_char_p,  # template name or NULL
+            c.c_int,  # max_tokens
+        ]
+        + sampling
+        + [TOKEN_CB, c.c_void_p]
+    )
+    shim.shim_chat.restype = c.c_int
+
+    shim.shim_complete.argtypes = (
+        [c.c_void_p, c.c_char_p, c.c_int] + sampling + [TOKEN_CB, c.c_void_p]
+    )
+    shim.shim_complete.restype = c.c_int
+
+    shim.shim_free.argtypes = [c.c_void_p]
+    shim.shim_free.restype = None
+
+
+def load() -> ctypes.CDLL:
+    """Load (once) and return the shim library."""
+    global _shim, _lib_dir
+    with _lock:
+        if _shim is not None:
+            return _shim
+
+        d = find_lib_dir()
+        try:
+            for name, required in _DEPENDENCIES:
+                path = d / name
+                if path.is_file():
+                    ctypes.CDLL(str(path), mode=ctypes.RTLD_GLOBAL)
+                elif required:
+                    raise LibraryNotFoundError("%s is missing from %s" % (name, d))
+            shim = ctypes.CDLL(str(d / SHIM_NAME))
+        except OSError as exc:
+            raise LibraryNotFoundError("failed to load native libraries from %s: %s" % (d, exc))
+
+        _bind(shim)
+        abi = shim.shim_abi_version()
+        if abi != SHIM_ABI_VERSION:
+            raise LibraryNotFoundError(
+                "native shim ABI %d does not match pylolm (expects %d); rebuild it with "
+                "`pylolm build-native`" % (abi, SHIM_ABI_VERSION)
+            )
+
+        shim.shim_set_verbose(0)
+        shim.shim_init(str(d).encode())
+        _shim = shim
+        _lib_dir = d
+        return shim
+
+
+def loaded_dir():
+    """Directory the libraries were loaded from, or None if not loaded yet."""
+    return _lib_dir

pylolm/errors.py

@@ -0,0 +1,25 @@
+"""Exceptions raised by pylolm."""
+
+
+class PyLoLmError(Exception):
+    """Base class for all pylolm errors."""
+
+
+class LibraryNotFoundError(PyLoLmError):
+    """The native libraries (libllama / libpylolm_shim) could not be found or loaded."""
+
+
+class ModelLoadError(PyLoLmError):
+    """The GGUF model could not be loaded."""
+
+
+class ContextOverflowError(PyLoLmError):
+    """The prompt does not fit in the context window (n_ctx)."""
+
+
+class GenerationError(PyLoLmError):
+    """llama.cpp failed while generating."""
+
+
+class DownloadError(PyLoLmError):
+    """Downloading or installing the prebuilt native libraries failed."""