diff --git a/ggml.py b/ggml.py new file mode 100644 index 0000000..1ea8d66 --- /dev/null +++ b/ggml.py @@ -0,0 +1,760 @@ +"""Speech to text and cleanup on this machine: whisper.cpp and llama.cpp. + +Two programs, one treatment. Fetch a release from GitHub, unpack it under the +data directory, fetch a model from Hugging Face, then keep one server alive on a +port of its own. Both of them speak the shape api.py already sends to the hosted +providers, so what the rest of Dikte sees is a base URL and nothing else: +whisper-server is started on `--inference-path /v1/audio/transcriptions`, the +exact path api.py builds, and llama-server answers /v1/chat/completions the way +OpenRouter does. + +A server rather than a one-shot run, because the model is the slow part. Loading +a large whisper model takes a second or two while transcribing a few seconds of +speech takes a fraction of one, and an LLM is worse: a server pays that once and +a run per dictation pays it every time. + +Nothing downloaded is trusted for having arrived. Every file is checked against +the sha256 its index published, and the bytes go to a `.part` that is only +renamed once the whole thing is there, so an interrupted download can never be +mistaken for a working one. + +This module imports hub and the string table, and nothing else of Dikte's: it +knows how to fetch a file and how to run a process, and nothing about dictation. +Its errors leave as LocalError and api.py turns them into the ApiError the +interface already knows how to show. +""" + +import atexit +import collections +import ctypes.util +import hashlib +import http.client +import json +import os +import pathlib +import platform +import shutil +import signal +import socket +import subprocess +import tarfile +import threading +import time +import urllib.error +import urllib.request + +import hub +from i18n import t + +HOST = "127.0.0.1" +# The path api.py asks for, so its URL and the server's line up. +INFERENCE_PATH = "/v1/audio/transcriptions" + +DATA_DIR = (pathlib.Path(os.environ.get("XDG_DATA_HOME") + or os.path.expanduser("~/.local/share")) / "dikte") +BIN_DIR = DATA_DIR / "bin" +MODELS_DIR = DATA_DIR / "models" + +# Loading a large model onto a GPU is the slow part of a start, and on a cold +# page cache a large LLM read from a spinning disk is slower still. +STARTUP_TIMEOUT = 180.0 +DOWNLOAD_CHUNK = 1 << 20 + +# `health` is the path that answers only once the model is in memory. whisper +# does not have one and does not need one: it binds its port after the model is +# loaded, so the port opening is the signal. +Program = collections.namedtuple("Program", "name repo binary health") + +WHISPER = Program("whisper", "ggml-org/whisper.cpp", "whisper-server", "") +LLAMA = Program("llama", "ggml-org/llama.cpp", "llama-server", "/health") + +# Where the models are listed. Neither list is written into Dikte: a catalogue +# in the source means a release of Dikte for every model somebody else +# publishes. +WHISPER_MODELS_REPO = "ggerganov/whisper.cpp" +LLM_AUTHOR = "ggml-org" + +# What the whisper repository holds besides models: Core ML encoders for Apple +# hardware and the odd loose file. +WHISPER_PREFIX = "ggml-" +WHISPER_SUFFIX = ".bin" + +# What a GGUF repository holds besides the model: mmproj is the vision half of a +# multimodal model, mtp a draft head for speculative decoding. Neither is a model +# a server can be started on, and offering them is offering a failure. +GGUF_SKIP = ("mmproj", "mtp-") +# Big enough for a 12B at Q4 and far past anything cleanup wants; the point is +# to keep a 400 GB frontier model out of a list somebody might click. +GGUF_MAX_BYTES = 16 << 30 + +# Suggestions, not a catalogue: the list itself is fetched, and these are only +# the rows that float to the top of it. Small instruction-following models, +# because cleanup is punctuation and filler words rather than anything that +# wants thinking about. +SUGGESTED_LLM = ( + "ggml-org/gemma-3-4b-it-GGUF", + "ggml-org/gemma-4-E2B-it-GGUF", + "ggml-org/gemma-4-E4B-it-GGUF", + "ggml-org/SmolLM3-3B-GGUF", +) +# Turbo at q5_0 is smaller than `small` and better than it, which makes the +# usual "start small" advice point at the same file as "start good". +SUGGESTED_WHISPER = "ggml-large-v3-turbo-q5_0.bin" + + +class LocalError(Exception): + pass + + +def human_size(count): + for unit in ("B", "KB", "MB", "GB"): + if count < 1024 or unit == "GB": + return f"{count:.0f} {unit}" if unit == "B" else f"{count:.1f} {unit}" + count /= 1024.0 + return f"{count:.1f} GB" + + +# --- fetching ------------------------------------------------------------- + + +def download(item, target, on_progress=None, should_stop=None): + """Fetch one hub.Item to `target`. True when it landed, False when stopped. + + The bytes go to a `.part` that is renamed only after both the length and the + hash agree with what the index said. A truncated file would otherwise sit + there looking installed and fail much later, inside a server, as a corrupt + model; a file that is the right length but the wrong content is worse, and + this is a program as often as it is a model. + """ + target = pathlib.Path(target) + part = target.with_name(target.name + ".part") + try: + target.parent.mkdir(parents=True, exist_ok=True) + except OSError as exc: + raise LocalError(t("Could not create {path}: {error}", + path=target.parent, error=exc)) from exc + + request = urllib.request.Request(item.url, headers={"User-Agent": hub.USER_AGENT}) + digest = hashlib.sha256() + done = 0 + try: + with urllib.request.urlopen(request, timeout=60) as response: + total = int(response.headers.get("Content-Length") or item.size or 0) + with open(part, "wb") as out: + while True: + if should_stop is not None and should_stop(): + part.unlink(missing_ok=True) + return False + block = response.read(DOWNLOAD_CHUNK) + if not block: + break + out.write(block) + digest.update(block) + done += len(block) + # More than was announced: a body that does not end is the + # one way this loop could run until the disk is full. + if total and done > total: + part.unlink(missing_ok=True) + raise LocalError(t("{name} is longer than it said it " + "would be.", name=item.name)) + if on_progress is not None: + on_progress(done, total) + # A proxy notice or an error page that came back as 200 would otherwise + # be renamed into place and only fail when something tries to read it. + if total and done != total: + part.unlink(missing_ok=True) + raise LocalError(t("The download stopped early ({done} of {total}).", + done=human_size(done), total=human_size(total))) + if item.sha256 and digest.hexdigest() != item.sha256: + part.unlink(missing_ok=True) + raise LocalError(t("{name} does not match its published checksum. " + "Nothing was installed.", name=item.name)) + part.replace(target) + return True + except urllib.error.HTTPError as exc: + part.unlink(missing_ok=True) + exc.close() # it holds the response body open until it is collected + raise LocalError(t("Could not download {name}: HTTP {code}", + name=item.name, code=exc.code)) from exc + except urllib.error.URLError as exc: + part.unlink(missing_ok=True) + raise LocalError(t("Could not download {name}: {error}", + name=item.name, error=exc.reason)) from exc + except OSError as exc: + # A connection cut mid-body arrives here too, and gigabytes in is + # exactly where that happens. + part.unlink(missing_ok=True) + raise LocalError(t("Could not write {name}: {error}", + name=item.name, error=exc)) from exc + + +# --- the programs --------------------------------------------------------- + + +def _arch(): + machine = platform.machine().lower() + if machine in ("aarch64", "arm64"): + return "arm64" + return "x64" + + +def _has_vulkan(): + """Whether a Vulkan loader is installed, which decides which build to fetch. + + llama.cpp publishes no CUDA build for Linux, so Vulkan is what a graphics + card gets here. The build without it is smaller and runs on the CPU, and + fetching the Vulkan one for a machine that cannot load it would only make + the download bigger. + """ + return bool(ctypes.util.find_library("vulkan")) + + +def _wanted_assets(program): + """Asset name endings to accept, best first.""" + arch = _arch() + if program is LLAMA and _has_vulkan(): + return (f"bin-ubuntu-vulkan-{arch}.tar.gz", f"bin-ubuntu-{arch}.tar.gz") + return (f"bin-ubuntu-{arch}.tar.gz",) + + +def _install_record(program): + return BIN_DIR / program.name / "installed.json" + + +def installed_program(program): + """The binary Dikte downloaded, or "" when there is none that still runs.""" + try: + record = json.loads(_install_record(program).read_text(encoding="utf-8")) + path = record.get("binary") or "" + except (OSError, ValueError): + return "" + return path if os.path.isfile(path) and os.access(path, os.X_OK) else "" + + +def installed_version(program): + try: + record = json.loads(_install_record(program).read_text(encoding="utf-8")) + return record.get("tag") or "" + except (OSError, ValueError): + return "" + + +def program_path(program, custom=""): + """Which copy of the program to run, or "" when there is none. + + A system one wins over a downloaded one. The distribution package is built + against whatever the machine has, which on this platform means it may reach + the graphics card, while the release binaries carry CPU backends only. + """ + custom = (custom or "").strip() + if custom: + return custom if os.path.isfile(custom) and os.access(custom, os.X_OK) else "" + return shutil.which(program.binary) or installed_program(program) + + +def system_program(program): + """Whether the program came from the system rather than from Dikte.""" + return bool(shutil.which(program.binary)) + + +def _find_binary(root, name): + for path in sorted(pathlib.Path(root).rglob(name)): + if path.is_file(): + return path + return None + + +def _extract(archive, into): + """Unpack a release tarball, refusing anything that reaches outside `into`. + + The archives lay their libraries next to their binaries and are linked with + an $ORIGIN runpath, so a whole directory is what has to survive the trip and + the binary cannot be lifted out of it. + """ + try: + with tarfile.open(archive, "r:gz") as tar: + try: + tar.extractall(into, filter="data") + except TypeError: # Python without the extraction filters + tar.extractall(into) + except (tarfile.TarError, OSError) as exc: + raise LocalError(t("Could not unpack {name}: {error}", + name=os.path.basename(str(archive)), error=exc)) from exc + + +def install_program(program, tag="", on_progress=None, should_stop=None, + refresh=False): + """Fetch and unpack a release. The path to the binary, or "" when stopped. + + `tag` is empty for whatever the project released last, which is the point: + a version pinned in Dikte's source would mean a release of Dikte every time + whisper.cpp has one. + """ + try: + tag, assets = hub.release(program.repo, tag or "latest", refresh=refresh) + except hub.HubError as exc: + raise LocalError(str(exc)) from exc + + item = None + for ending in _wanted_assets(program): + item = next((a for a in assets if a.name.endswith(ending)), None) + if item: + break + if item is None: + raise LocalError(t("{repo} {tag} has no build for this machine.", + repo=program.repo, tag=tag)) + + into = BIN_DIR / program.name / tag + shutil.rmtree(into, ignore_errors=True) + archive = BIN_DIR / program.name / item.name + try: + if not download(item, archive, on_progress, should_stop): + return "" + _extract(archive, into) + binary = _find_binary(into, program.binary) + if binary is None: + raise LocalError(t("{name} was not in the download.", + name=program.binary)) + binary.chmod(binary.stat().st_mode | 0o111) + _install_record(program).write_text( + json.dumps({"tag": tag, "binary": str(binary)}), encoding="utf-8") + except OSError as exc: + raise LocalError(t("Could not install {name}: {error}", + name=program.name, error=exc)) from exc + finally: + try: + archive.unlink(missing_ok=True) + except OSError: + pass + _drop_old_versions(program, keep=tag) + return str(binary) + + +def _drop_old_versions(program, keep): + """Leave one unpacked release behind, not one per update.""" + root = BIN_DIR / program.name + try: + for path in root.iterdir(): + if path.is_dir() and path.name != keep: + shutil.rmtree(path, ignore_errors=True) + except OSError: + pass + + +# --- the models ----------------------------------------------------------- + + +def whisper_models(refresh=False): + """[hub.Item] for every whisper model on offer, smallest first.""" + try: + files = hub.files(WHISPER_MODELS_REPO, refresh=refresh) + except hub.HubError as exc: + raise LocalError(str(exc)) from exc + models = [f for f in files + if f.name.startswith(WHISPER_PREFIX) and f.name.endswith(WHISPER_SUFFIX) + and f.size > 0] + return sorted(models, key=lambda f: f.size) + + +def llm_repos(refresh=False): + """Repository ids for the GGUF models on offer, suggestions first.""" + try: + found = [r.id for r in hub.repos(author=LLM_AUTHOR, refresh=refresh)] + except hub.HubError: + # A menu rather than a catalogue: with nothing to show, the suggestions + # are still worth showing, and whatever is wrong with the network will + # say so where it matters, when a download is asked for. + found = [] + if not found: + return list(SUGGESTED_LLM) + first = [r for r in SUGGESTED_LLM if r in found] + return first + [r for r in found if r not in first] + + +def llm_quants(repo, refresh=False): + """[hub.Item] for the model files in one GGUF repository, smallest first.""" + try: + files = hub.files(repo, refresh=refresh) + except hub.HubError as exc: + raise LocalError(str(exc)) from exc + out = [] + for item in files: + name = item.name.rsplit("/", 1)[-1] + if not name.endswith(".gguf") or name.startswith(GGUF_SKIP): + continue + # A model split across files needs all of them and a different command + # line; anything cleanup wants fits in one. + if "-of-000" in name or not 0 < item.size <= GGUF_MAX_BYTES: + continue + out.append(item) + return sorted(out, key=lambda f: f.size) + + +def whisper_model_path(name): + return MODELS_DIR / "whisper" / name + + +def llm_model_path(name): + return MODELS_DIR / "llm" / name.rsplit("/", 1)[-1] + + +def have_model(path): + path = pathlib.Path(path) + return path.is_file() and path.stat().st_size > 0 + + +def installed_whisper_models(): + return sorted(p.name for p in (MODELS_DIR / "whisper").glob("*.bin")) + + +def installed_llm_models(): + return sorted(p.name for p in (MODELS_DIR / "llm").glob("*.gguf")) + + +def delete_model(path): + try: + pathlib.Path(path).unlink() + except FileNotFoundError: + pass + except OSError as exc: + raise LocalError(t("Could not delete the model: {error}", error=exc)) from exc + + +# --- one server ----------------------------------------------------------- + + +def _free_port(): + """A port nothing is listening on, handed straight to the server. + + Between closing this socket and the server binding it, something else could + take it; that is why a start retries rather than trusting the number. + """ + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind((HOST, 0)) + return sock.getsockname()[1] + + +def _listening(port): + try: + with socket.create_connection((HOST, port), timeout=0.5): + return True + except OSError: + return False + + +def _healthy(port, path): + """Whether the model is in memory, for a server that says so. + + Spoken over http.client rather than urllib because this never leaves the + machine: it is the same question as _listening, one layer up. + """ + connection = http.client.HTTPConnection(HOST, port, timeout=2) + try: + connection.request("GET", path) + # 503 for as long as the model is still being read in. + return connection.getresponse().status == 200 + except (http.client.HTTPException, OSError): + return False + finally: + connection.close() + + +def _tail(path, lines=3): + try: + with open(path, encoding="utf-8", errors="replace") as fh: + found = [line.strip() for line in fh if line.strip()] + except OSError: + return "" + return " | ".join(found[-lines:]) + + +class Server: + """One process, started when something needs it and stopped when nothing does. + + `build` turns the settings into a command line; everything else about + running a server is the same for both programs. + """ + + def __init__(self, program, build, defaults): + self.program = program + self._build = build + self._settings = dict(defaults) + # Two locks on purpose. `_lock` is held for the length of a dictionary + # lookup, so the interface can ask what is running while a model is + # being loaded; `_starting` is held across the start itself, which can + # take a minute and which two threads must not both do. + self._lock = threading.Lock() + self._starting = threading.Lock() + self._proc = None + self._port = 0 + self._log = "" + self._key = None + + # ---- settings -------------------------------------------------------- + + def configure(self, **changes): + """Apply settings. A server started on the old ones is stopped.""" + with self._lock: + for key, value in changes.items(): + if value is not None and key in self._settings: + self._settings[key] = value + stale = self._proc is not None and self._key != self._settings_key() + if stale: + self.stop() + + def settings(self): + with self._lock: + return dict(self._settings) + + def _settings_key(self): + """What a running server would have to be restarted for.""" + return json.dumps(self._settings, sort_keys=True, default=str) + + # ---- process --------------------------------------------------------- + + @property + def running(self): + with self._lock: + return self._proc is not None and self._proc.poll() is None + + def base_url(self): + with self._lock: + return f"http://{HOST}:{self._port}/v1" if self._port else "" + + def error(self): + """The last thing the server printed, for a failure after it started.""" + with self._lock: + log = self._log + return _tail(log) if log else "" + + def serve(self): + """The base URL of a server that is up and running the current settings.""" + ready = self._current_url() + if ready: + return ready + with self._starting: + # Somebody may have started it while this thread waited its turn. + ready = self._current_url() + if ready: + return ready + self.stop() + with self._lock: + settings, key = dict(self._settings), self._settings_key() + proc, port, log = self._launch(settings) + with self._lock: + self._proc, self._port, self._log, self._key = proc, port, log, key + return self.base_url() + + def _current_url(self): + with self._lock: + up = self._proc is not None and self._proc.poll() is None + return (f"http://{HOST}:{self._port}/v1" + if up and self._key == self._settings_key() else "") + + def _launch(self, settings): + args = self._build(settings) # raises LocalError when unusable + last = "" + for _ in range(3): + port = _free_port() + log = DATA_DIR / f"{self.program.name}-server.log" + try: + log.parent.mkdir(parents=True, exist_ok=True) + sink = open(log, "wb") + except OSError as exc: + raise LocalError(t("Could not start {name}: {error}", + name=self.program.name, error=exc)) from exc + try: + with sink: + proc = subprocess.Popen( + args + ["--host", HOST, "--port", str(port)], + stdout=sink, stderr=subprocess.STDOUT, + stdin=subprocess.DEVNULL, + ) + except OSError as exc: + raise LocalError(t("Could not start {name}: {error}", + name=self.program.name, error=exc)) from exc + + # Written before it is ready rather than after, so that a kill + # during the model load leaves something for the sweep to find. + self._remember(proc.pid) + if self._wait_ready(proc, port): + return proc, port, str(log) + last = _tail(log) + self._forget() + # A port taken between the probe and the bind is the one failure + # worth another go; anything else will fail the same way again. + if "address" not in last.lower() and "bind" not in last.lower(): + break + raise LocalError(t("{name} did not start: {error}", + name=self.program.binary, error=last or t("no output"))) + + def _wait_ready(self, proc, port): + deadline = time.monotonic() + STARTUP_TIMEOUT + while time.monotonic() < deadline: + if proc.poll() is not None: + return False + if _listening(port): + # whisper binds after the model is loaded, so the open port is + # the answer. llama binds first and answers /health with 503 + # until it is ready. + if not self.program.health or _healthy(port, self.program.health): + return True + time.sleep(0.1) + proc.kill() + proc.wait(timeout=5) + return False + + def stop(self): + with self._lock: + proc, self._proc = self._proc, None + self._port, self._log, self._key = 0, "", None + if proc is not None and proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=5) + except subprocess.TimeoutExpired: + proc.kill() + proc.wait(timeout=5) + if proc is not None: + self._forget() + + # ---- servers a killed Dikte left behind ------------------------------- + + def _pid_file(self): + return DATA_DIR / f"{self.program.name}-server.pid" + + def _remember(self, pid): + try: + path = self._pid_file() + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(str(pid)) + except OSError: + pass # the sweep is a safety net, not something to fail a run over + + def _forget(self): + try: + self._pid_file().unlink() + except OSError: + pass + + def _is_ours(self, pid): + """Whether that pid is still the server this Dikte started. + + Asked because pids are handed out again: by the time anyone looks, the + number could belong to something else entirely, and killing it would be + a good deal worse than the leak being cleaned up. The program name alone + could be somebody else's copy; the name together with Dikte's own data + directory on the command line could not. + """ + try: + blob = pathlib.Path(f"/proc/{pid}/cmdline").read_bytes() + except OSError: + return False + return (self.program.binary.encode() in blob + and str(DATA_DIR).encode() in blob) + + def sweep(self): + """Kill a server a previous Dikte left behind. True when one was found. + + stop() and atexit cover every exit that gets to run code. A SIGKILL does + not, and neither does a session torn down from under it, and the server + would then sit there holding the model with nothing left alive to ask it + anything. + """ + try: + pid = int(self._pid_file().read_text().strip()) + except (OSError, ValueError): + return False + self._forget() + if not self._is_ours(pid): + return False + try: + os.kill(pid, signal.SIGTERM) + except OSError: + return False + return True + + +# --- the two of them ------------------------------------------------------ + + +def _whisper_args(settings): + binary = program_path(WHISPER, settings["binary"]) + if not binary: + raise LocalError(t("whisper.cpp is not installed. Settings → API and " + "models → Download.")) + model = whisper_model_path(settings["model"]) + if not settings["model"] or not have_model(model): + raise LocalError(t("No whisper model has been downloaded yet. " + "Settings → API and models → Download.")) + args = [ + binary, "-m", str(model), + "--inference-path", INFERENCE_PATH, + # Whatever language the request does not name. api.py leaves the field + # out when the language is "auto", and the server's own default is + # English rather than detection. + "-l", "auto", + # Stock phrases invented for near-silence come from non-speech tokens, + # and verbose_json otherwise pays for a language probability sweep + # nothing here reads. + "-sns", "-nlp", + ] + if int(settings["threads"]) > 0: + args += ["-t", str(int(settings["threads"]))] + if not settings["gpu"]: + args.append("-ng") + return args + + +def _llm_args(settings): + binary = program_path(LLAMA, settings["binary"]) + if not binary: + raise LocalError(t("llama.cpp is not installed. Settings → API and " + "models → Download.")) + model = llm_model_path(settings["model"]) + if not settings["model"] or not have_model(model): + raise LocalError(t("No local cleanup model has been downloaded yet. " + "Settings → API and models → Download.")) + args = [binary, "-m", str(model), "-c", str(int(settings["context"]))] + # All of them, or as many as fit: llama.cpp stops offloading when the card + # is full rather than failing, and a build with no GPU backend ignores it. + args += ["-ngl", "99" if settings["gpu"] else "0"] + if int(settings["threads"]) > 0: + args += ["-t", str(int(settings["threads"]))] + return args + + +whisper = Server(WHISPER, _whisper_args, { + "model": "", + "threads": 0, + "gpu": True, + "binary": "", +}) + +llm = Server(LLAMA, _llm_args, { + "model": "", + "threads": 0, + "gpu": True, + "binary": "", + # A dictation and its prompt are short. This is sized for the longest + # cleanup block rather than for a conversation, and it is what the model + # costs in memory beyond its own weights. + "context": 8192, +}) + +SERVERS = (whisper, llm) + + +def sweep(): + """Clean up after a Dikte that was killed outright. True when one was found.""" + return any([server.sweep() for server in SERVERS]) + + +def stop_all(): + for server in SERVERS: + server.stop() + + +# Dikte stops the servers itself on quit and on restart; this catches the paths +# that skip that, such as an unhandled exception on the way out. +atexit.register(stop_all) diff --git a/hub.py b/hub.py new file mode 100644 index 0000000..f793d1c --- /dev/null +++ b/hub.py @@ -0,0 +1,186 @@ +"""Where the programs and the models come from: GitHub releases and Hugging Face. + +Both answer plain JSON over HTTPS without a key, and both publish a sha256 for +every file they hand out: GitHub as the asset digest, Hugging Face as the LFS +object id. Nothing that lands on disk is trusted for having arrived, which +matters more here than it usually would, because half of what is fetched is a +program Dikte then runs. + +The lists are read rather than kept. A model catalogue written into the source +means a release of Dikte for every new model, and a pinned whisper.cpp version +means one for every whisper.cpp release; both of those are somebody else's news, +not Dikte's. Answers are cached for a few hours, and a cache that has gone stale +is still a better answer than none when the network is down. + +Nothing here imports the rest of Dikte apart from the string table: this module +knows two websites and nothing about dictation. +""" + +import collections +import json +import os +import pathlib +import time +import urllib.error +import urllib.parse +import urllib.request + +from i18n import t + +GITHUB_API = "https://api.github.com" +HF_API = "https://huggingface.co/api" +HF_FILES = "https://huggingface.co" +USER_AGENT = "dikte/1.0 (+https://github.com/yusufipk/dikte)" + +CACHE_DIR = (pathlib.Path(os.environ.get("XDG_CACHE_HOME") + or os.path.expanduser("~/.cache")) / "dikte") +# Long enough that opening the settings window twice in an evening asks nobody +# anything, short enough that a model published this morning is offered today. +CACHE_TTL = 6 * 3600 + +# `sha256` is empty for the few files neither side stores in LFS; those are the +# small ones, and a checksum is only worth having where there is something to +# check. +Item = collections.namedtuple("Item", "name url size sha256") +Repo = collections.namedtuple("Repo", "id downloads updated") + + +class HubError(Exception): + pass + + +def _get(url, timeout=20): + request = urllib.request.Request(url, headers={ + "User-Agent": USER_AGENT, + "Accept": "application/json", + }) + try: + with urllib.request.urlopen(request, timeout=timeout) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + exc.close() # it holds the response body open until it is collected + raise HubError(t("{url} answered HTTP {code}.", + url=urllib.parse.urlsplit(url).netloc, code=exc.code)) from exc + except urllib.error.URLError as exc: + raise HubError(t("Could not reach {url}: {error}", + url=urllib.parse.urlsplit(url).netloc, + error=exc.reason)) from exc + except (ValueError, OSError) as exc: + raise HubError(t("Could not read the answer from {url}: {error}", + url=urllib.parse.urlsplit(url).netloc, error=exc)) from exc + + +def _cache_file(key): + safe = "".join(c if c.isalnum() or c in "-._" else "-" for c in key) + return CACHE_DIR / f"{safe}.json" + + +def _read_cache(key, ttl): + """What was stored under this key, or None. `ttl` of 0 ignores the age.""" + path = _cache_file(key) + try: + age = time.time() - path.stat().st_mtime + if ttl and age > ttl: + return None + return json.loads(path.read_text(encoding="utf-8")) + except (OSError, ValueError): + return None + + +def _write_cache(key, payload): + try: + CACHE_DIR.mkdir(parents=True, exist_ok=True) + _cache_file(key).write_text(json.dumps(payload), encoding="utf-8") + except OSError: + pass # a cache that cannot be written is not a failed lookup + + +def _fetch(key, url, ttl=CACHE_TTL, refresh=False): + """The JSON at `url`, from the cache when it is fresh enough. + + A lookup that fails falls back to the cache however old it is: an offline + settings window that shows yesterday's list is worth a great deal more than + one that shows an error. + """ + if not refresh: + cached = _read_cache(key, ttl) + if cached is not None: + return cached + try: + payload = _get(url) + except HubError: + stale = _read_cache(key, 0) + if stale is not None: + return stale + raise + _write_cache(key, payload) + return payload + + +def _digest(value): + """GitHub writes its digests as "sha256:…"; Hugging Face writes the hash.""" + value = (value or "").strip() + return value.split(":", 1)[1] if value.startswith("sha256:") else value + + +def release(repo, tag="latest", refresh=False): + """(tag, [Item]) for one GitHub release, newest when no tag is given.""" + where = "latest" if tag in ("", "latest") else f"tags/{tag}" + data = _fetch(f"gh-{repo}-{tag or 'latest'}", + f"{GITHUB_API}/repos/{repo}/releases/{where}", refresh=refresh) + if not isinstance(data, dict) or not data.get("assets"): + raise HubError(t("{repo} has no downloadable release.", repo=repo)) + assets = [Item(a.get("name") or "", a.get("browser_download_url") or "", + int(a.get("size") or 0), _digest(a.get("digest"))) + for a in data["assets"] if a.get("browser_download_url")] + return data.get("tag_name") or tag, assets + + +def files(repo, revision="main", refresh=False): + """[Item] for every file in a Hugging Face repository. + + The size is there whether or not the file is in LFS; the hash is only there + when it is, which for anything worth downloading it always is. + """ + data = _fetch(f"hf-tree-{repo}-{revision}", + f"{HF_API}/models/{repo}/tree/{revision}?recursive=true", + refresh=refresh) + if not isinstance(data, list): + raise HubError(t("{repo} did not return a file list.", repo=repo)) + out = [] + for entry in data: + if entry.get("type") != "file": + continue + path = entry.get("path") or "" + lfs = entry.get("lfs") or {} + out.append(Item( + path, + f"{HF_FILES}/{repo}/resolve/{revision}/{urllib.parse.quote(path)}", + int(lfs.get("size") or entry.get("size") or 0), + _digest(lfs.get("oid") or lfs.get("sha256")), + )) + return out + + +def repos(author="", search="", limit=40, refresh=False): + """[Repo] of GGUF repositories, newest first. + + Filtered by author on purpose. Hugging Face's own trending list is open to + everyone and reads like it: asking it for the popular GGUF today answers + with a wall of roleplay merges, which is not what a dictation transcript + wants cleaning up. An author is a small enough thing to trust and a large + enough one to keep the list current without Dikte being updated. + """ + query = {"filter": "gguf", "sort": "lastModified", "direction": "-1", + "limit": str(limit)} + if author: + query["author"] = author + if search: + query["search"] = search + url = f"{HF_API}/models?{urllib.parse.urlencode(query)}" + data = _fetch(f"hf-models-{author}-{search}-{limit}", url, refresh=refresh) + if not isinstance(data, list): + raise HubError(t("Hugging Face did not return a model list.")) + return [Repo(m.get("id") or "", int(m.get("downloads") or 0), + m.get("lastModified") or "") + for m in data if m.get("id")] diff --git a/tests/test_ggml.py b/tests/test_ggml.py new file mode 100644 index 0000000..ea2af74 --- /dev/null +++ b/tests/test_ggml.py @@ -0,0 +1,582 @@ +"""Fetching a program and a model, and keeping a server alive on them. + +No network and no whisper.cpp: the downloads are answered from memory, and the +servers are stand-in scripts that take the same arguments and open their port +when they are told to, which is the only thing the code waits on. +""" + +import contextlib +import hashlib +import io +import os +import signal +import sys +import tarfile +import textwrap +import threading +import time +from unittest import mock + +import ggml +import hub +from tests.support import (DikteTest, fake_urlopen, http_error, json_body, + linux_only, url_error) + + +def body(data, length=None): + """What urlopen hands back for a download: a reader with a length header.""" + class Body: + def __init__(self): + self._buf = io.BytesIO(data) + self.headers = {"Content-Length": + str(len(data) if length is None else length)} + + def read(self, count=-1): + return self._buf.read(count) + + def __enter__(self): + return self + + def __exit__(self, *_): + return False + return Body() + + +def item(name, data, url="https://example.invalid/f", sha=True): + return hub.Item(name, url, len(data), + hashlib.sha256(data).hexdigest() if sha else "") + + +@contextlib.contextmanager +def serving(release, archive): + """Answer by what is being asked for rather than by what came before. + + An install asks GitHub what the release is and then asks for one file out of + it, and the first of those two comes from the cache the second time around. + Answering in order would then hand the archive request the release listing. + """ + def opener(request, timeout=None): + url = request.full_url + if "api.github.com" in url: + return json_body(release) + return body(archive) + + with mock.patch("urllib.request.urlopen", side_effect=opener) as calls: + yield calls + + +def tarball(entries): + """A .tar.gz laid out the way the releases are: one directory of files.""" + buf = io.BytesIO() + with tarfile.open(fileobj=buf, mode="w:gz") as tar: + for name, content in entries.items(): + info = tarfile.TarInfo(name) + info.size = len(content) + info.mode = 0o755 + tar.addfile(info, io.BytesIO(content)) + return buf.getvalue() + + +class Local(DikteTest): + """A test with its own bin, models and cache directories.""" + + def setUp(self): + super().setUp() + self.patch_attr(ggml, "DATA_DIR", self.path("data")) + self.patch_attr(ggml, "BIN_DIR", self.path("data", "bin")) + self.patch_attr(ggml, "MODELS_DIR", self.path("data", "models")) + self.patch_attr(hub, "CACHE_DIR", self.path("cache")) + + +# --- downloading ---------------------------------------------------------- + + +class Download(Local): + def test_it_lands_and_the_part_file_is_gone(self): + data = b"a model, more or less" * 100 + target = self.path("data", "models", "m.bin") + with fake_urlopen(body(data)): + self.assertTrue(ggml.download(item("m.bin", data), target)) + self.assertEqual(target.read_bytes(), data) + self.assertFalse(target.with_name("m.bin.part").exists()) + + def test_a_wrong_checksum_installs_nothing(self): + data = b"the bytes that arrived" + wrong = hub.Item("m.bin", "https://example.invalid/f", len(data), "f" * 64) + target = self.path("data", "models", "m.bin") + with fake_urlopen(body(data)): + with self.assertRaises(ggml.LocalError) as caught: + ggml.download(wrong, target) + self.assertIn("checksum", str(caught.exception)) + self.assertFalse(target.exists()) + self.assertFalse(target.with_name("m.bin.part").exists()) + + def test_a_body_shorter_than_its_header_installs_nothing(self): + data = b"half of it" + target = self.path("data", "models", "m.bin") + with fake_urlopen(body(data, length=len(data) * 2)): + with self.assertRaises(ggml.LocalError): + ggml.download(item("m.bin", data), target) + self.assertFalse(target.exists()) + + def test_a_file_with_no_published_checksum_is_still_taken(self): + data = b"a README, say" + target = self.path("data", "models", "readme") + with fake_urlopen(body(data)): + self.assertTrue(ggml.download(item("readme", data, sha=False), target)) + self.assertTrue(target.exists()) + + def test_stopping_leaves_nothing_behind(self): + data = b"x" * (ggml.DOWNLOAD_CHUNK * 3) + target = self.path("data", "models", "m.bin") + with fake_urlopen(body(data)): + landed = ggml.download(item("m.bin", data), target, + should_stop=lambda: True) + self.assertFalse(landed) + self.assertFalse(target.exists()) + self.assertFalse(target.with_name("m.bin.part").exists()) + + def test_progress_is_reported_against_the_total(self): + data = b"y" * (ggml.DOWNLOAD_CHUNK + 5) + seen = [] + with fake_urlopen(body(data)): + ggml.download(item("m.bin", data), self.path("data", "m.bin"), + on_progress=lambda done, total: seen.append((done, total))) + self.assertEqual(seen[-1], (len(data), len(data))) + self.assertGreater(len(seen), 1) + + def test_a_refused_connection_says_which_file(self): + with fake_urlopen(url_error("no route to host")): + with self.assertRaises(ggml.LocalError) as caught: + ggml.download(item("m.bin", b"x"), self.path("data", "m.bin")) + self.assertIn("m.bin", str(caught.exception)) + + def test_an_http_error_is_not_written_to_disk(self): + target = self.path("data", "m.bin") + with fake_urlopen(http_error(404)): + with self.assertRaises(ggml.LocalError): + ggml.download(item("m.bin", b"x"), target) + self.assertFalse(target.exists()) + + +# --- installing a program ------------------------------------------------- + + +class InstallProgram(Local): + def release(self, *names): + return {"tag_name": "v1.9.1", "assets": [ + {"name": name, "browser_download_url": f"https://example.invalid/{name}", + "size": 10, "digest": ""} for name in names]} + + def archive(self): + return tarball({ + "whisper-bin-ubuntu-x64/whisper-server": b"#!/bin/sh\nexit 0\n", + "whisper-bin-ubuntu-x64/libwhisper.so": b"not really a library", + }) + + def install(self, *names, archive=None): + self.patch_attr(ggml, "_arch", lambda: "x64") + with serving(self.release(*names), + self.archive() if archive is None else archive) as calls: + path = ggml.install_program(ggml.WHISPER) + return path, [call.args[0].full_url for call in calls.call_args_list] + + def test_the_binary_and_its_libraries_land_together(self): + path, _ = self.install("whisper-bin-ubuntu-x64.tar.gz") + self.assertTrue(os.path.isfile(path)) + self.assertTrue(os.access(path, os.X_OK)) + self.assertTrue(os.path.isfile(os.path.join(os.path.dirname(path), + "libwhisper.so"))) + + def test_the_build_for_this_machine_is_the_one_fetched(self): + _, urls = self.install("whisper-bin-x64.zip", "whisper-bin-ubuntu-arm64.tar.gz", + "whisper-bin-ubuntu-x64.tar.gz") + self.assertTrue(urls[1].endswith("whisper-bin-ubuntu-x64.tar.gz")) + + def test_a_release_with_nothing_for_this_machine_says_so(self): + self.patch_attr(ggml, "_arch", lambda: "x64") + with fake_urlopen(self.release("whisper-bin-Win32.zip")): + with self.assertRaises(ggml.LocalError) as caught: + ggml.install_program(ggml.WHISPER) + self.assertIn("this machine", str(caught.exception)) + + def test_what_was_installed_is_remembered(self): + path, _ = self.install("whisper-bin-ubuntu-x64.tar.gz") + self.assertEqual(ggml.installed_program(ggml.WHISPER), path) + self.assertEqual(ggml.installed_version(ggml.WHISPER), "v1.9.1") + + def test_a_record_pointing_at_a_deleted_binary_counts_for_nothing(self): + path, _ = self.install("whisper-bin-ubuntu-x64.tar.gz") + os.unlink(path) + self.assertEqual(ggml.installed_program(ggml.WHISPER), "") + + def test_the_archive_is_not_kept(self): + self.install("whisper-bin-ubuntu-x64.tar.gz") + left = list((self.path("data", "bin", "whisper")).glob("*.tar.gz")) + self.assertEqual(left, []) + + def test_the_previous_version_is_swept_up(self): + self.install("whisper-bin-ubuntu-x64.tar.gz") + old = self.path("data", "bin", "whisper", "v1.9.0") + old.mkdir(parents=True) + (old / "whisper-server").write_bytes(b"older") + self.install("whisper-bin-ubuntu-x64.tar.gz") + self.assertFalse(old.exists()) + + def test_an_archive_without_the_binary_is_refused(self): + empty = tarball({"whisper-bin-ubuntu-x64/README": b"nothing here"}) + with self.assertRaises(ggml.LocalError) as caught: + self.install("whisper-bin-ubuntu-x64.tar.gz", archive=empty) + self.assertIn("whisper-server", str(caught.exception)) + + def test_llama_takes_the_vulkan_build_when_there_is_a_loader(self): + self.patch_attr(ggml, "_arch", lambda: "x64") + self.patch_attr(ggml, "_has_vulkan", lambda: True) + self.assertEqual(ggml._wanted_assets(ggml.LLAMA)[0], + "bin-ubuntu-vulkan-x64.tar.gz") + + def test_llama_falls_back_to_the_plain_build_without_one(self): + self.patch_attr(ggml, "_arch", lambda: "x64") + self.patch_attr(ggml, "_has_vulkan", lambda: False) + self.assertEqual(ggml._wanted_assets(ggml.LLAMA), ("bin-ubuntu-x64.tar.gz",)) + + +class WhichCopyRuns(Local): + def test_a_system_build_wins_over_a_downloaded_one(self): + self.patch_attr(ggml, "installed_program", lambda program: "/data/whisper-server") + with mock.patch("shutil.which", return_value="/usr/bin/whisper-server"): + self.assertEqual(ggml.program_path(ggml.WHISPER), "/usr/bin/whisper-server") + + def test_the_downloaded_one_is_used_when_there_is_no_system_build(self): + self.patch_attr(ggml, "installed_program", lambda program: "/data/whisper-server") + with mock.patch("shutil.which", return_value=None): + self.assertEqual(ggml.program_path(ggml.WHISPER), "/data/whisper-server") + + def test_a_setting_pointing_at_nothing_is_no_program(self): + self.assertEqual(ggml.program_path(ggml.WHISPER, "/nowhere/whisper-server"), "") + + def test_a_setting_pointing_at_a_program_wins(self): + mine = self.path("mine") + mine.write_text("#!/bin/sh\n") + mine.chmod(0o755) + with mock.patch("shutil.which", return_value="/usr/bin/whisper-server"): + self.assertEqual(ggml.program_path(ggml.WHISPER, str(mine)), str(mine)) + + +# --- the lists ------------------------------------------------------------ + + +WHISPER_TREE = [ + {"type": "file", "path": "ggml-base.bin", "size": 147951465, + "lfs": {"oid": "a" * 64}}, + {"type": "file", "path": "ggml-large-v3-turbo-q5_0.bin", "size": 574041195, + "lfs": {"oid": "b" * 64}}, + {"type": "file", "path": "ggml-base-encoder.mlmodelc.zip", "size": 37922638, + "lfs": {"oid": "c" * 64}}, + {"type": "file", "path": "README.md", "size": 3196}, +] + +GGUF_TREE = [ + {"type": "file", "path": "gemma-3-4b-it-Q4_K_M.gguf", "size": 2489000000, + "lfs": {"oid": "a" * 64}}, + {"type": "file", "path": "gemma-3-4b-it-Q8_0.gguf", "size": 4130000000, + "lfs": {"oid": "b" * 64}}, + {"type": "file", "path": "mmproj-model-f16.gguf", "size": 851000000, + "lfs": {"oid": "c" * 64}}, + {"type": "file", "path": "mtp-gemma-4-E4B-it-Q4_0.gguf", "size": 59000000, + "lfs": {"oid": "d" * 64}}, + {"type": "file", "path": "huge-00001-of-00009.gguf", "size": 40000000000, + "lfs": {"oid": "e" * 64}}, + {"type": "file", "path": "README.md", "size": 100}, +] + + +class Catalogue(Local): + def test_only_models_are_offered_and_the_small_ones_first(self): + with fake_urlopen(WHISPER_TREE): + models = ggml.whisper_models() + self.assertEqual([m.name for m in models], + ["ggml-base.bin", "ggml-large-v3-turbo-q5_0.bin"]) + + def test_the_core_ml_encoders_are_not_models(self): + with fake_urlopen(WHISPER_TREE): + names = [m.name for m in ggml.whisper_models()] + self.assertNotIn("ggml-base-encoder.mlmodelc.zip", names) + + def test_the_projector_and_the_draft_head_are_not_models(self): + with fake_urlopen(GGUF_TREE): + names = [q.name for q in ggml.llm_quants("ggml-org/gemma-3-4b-it-GGUF")] + self.assertEqual(names, + ["gemma-3-4b-it-Q4_K_M.gguf", "gemma-3-4b-it-Q8_0.gguf"]) + + def test_a_model_split_across_files_is_left_out(self): + with fake_urlopen(GGUF_TREE): + names = [q.name for q in ggml.llm_quants("ggml-org/gemma-3-4b-it-GGUF")] + self.assertNotIn("huge-00001-of-00009.gguf", names) + + def test_the_suggestions_come_first_and_the_rest_follow(self): + listing = [{"id": "ggml-org/something-new-GGUF"}, + {"id": ggml.SUGGESTED_LLM[0]}] + with fake_urlopen(listing): + found = ggml.llm_repos() + self.assertEqual(found[0], ggml.SUGGESTED_LLM[0]) + self.assertIn("ggml-org/something-new-GGUF", found) + + def test_an_unreachable_list_still_offers_the_suggestions(self): + with fake_urlopen(url_error()): + self.assertEqual(ggml.llm_repos(), list(ggml.SUGGESTED_LLM)) + + def test_an_unreachable_whisper_list_is_an_error_worth_showing(self): + with fake_urlopen(url_error()): + with self.assertRaises(ggml.LocalError): + ggml.whisper_models() + + def test_what_is_on_disk_is_read_from_disk(self): + self.assertEqual(ggml.installed_whisper_models(), []) + path = ggml.whisper_model_path("ggml-base.bin") + path.parent.mkdir(parents=True) + path.write_bytes(b"model") + self.assertEqual(ggml.installed_whisper_models(), ["ggml-base.bin"]) + self.assertTrue(ggml.have_model(path)) + + def test_an_empty_file_is_not_a_model(self): + path = ggml.llm_model_path("ggml-org/x-GGUF/model.gguf") + path.parent.mkdir(parents=True) + path.write_bytes(b"") + self.assertFalse(ggml.have_model(path)) + + def test_a_model_is_named_by_its_file_not_its_repository(self): + self.assertEqual(ggml.llm_model_path("ggml-org/x-GGUF/model.gguf").name, + "model.gguf") + + +# --- keeping a server alive ----------------------------------------------- + + +STAND_IN = textwrap.dedent(""" + import http.server, sys, threading, time + + args = sys.argv[1:] + + def opt(name, default=""): + return args[args.index(name) + 1] if name in args else default + + if "--die" in args: + print("could not load model: no such file") + sys.exit(2) + + time.sleep(float(opt("--wait", "0"))) + + started = time.monotonic() + healthy_after = float(opt("--healthy-after", "0")) + + class Handler(http.server.BaseHTTPRequestHandler): + def do_GET(self): + ok = time.monotonic() - started >= healthy_after + self.send_response(200 if ok else 503) + self.end_headers() + self.wfile.write(b"{}") + + def log_message(self, *a): + pass + + server = http.server.HTTPServer((opt("--host"), int(opt("--port"))), Handler) + print("listening on " + opt("--port"), flush=True) + server.serve_forever() +""") + + +class Servers(Local): + def setUp(self): + super().setUp() + self.path("data").mkdir(parents=True, exist_ok=True) + # Named for the program and kept inside the data directory, because that + # is what the sweep looks for on a command line. + self.script = self.path("data", "whisper-server.py") + self.script.write_text(STAND_IN) + self.addCleanup(ggml.stop_all) + self.servers = [] + + def server(self, program=ggml.WHISPER, **settings): + defaults = {"extra": []} + defaults.update(settings) + made = ggml.Server( + program, + lambda values: [sys.executable, str(self.script)] + list(values["extra"]), + defaults, + ) + self.servers.append(made) + self.addCleanup(made.stop) + return made + + def test_a_started_server_hands_back_its_address(self): + server = self.server() + url = server.serve() + self.assertRegex(url, r"^http://127\.0\.0\.1:\d+/v1$") + self.assertTrue(server.running) + + def test_the_second_call_does_not_start_a_second_one(self): + server = self.server() + first = server.serve() + self.assertEqual(server.serve(), first) + + def test_a_settings_change_stops_what_was_running(self): + server = self.server() + server.serve() + server.configure(extra=["--wait", "0"]) + self.assertFalse(server.running) + + def test_the_new_settings_are_what_the_next_start_uses(self): + server = self.server() + server.serve() + server.configure(extra=["--healthy-after", "0"]) + second = server.serve() + self.assertTrue(server.running) + self.assertTrue(second) + + def test_a_program_that_dies_reports_what_it_printed(self): + server = self.server(extra=["--die"]) + with self.assertRaises(ggml.LocalError) as caught: + server.serve() + self.assertIn("no such file", str(caught.exception)) + self.assertFalse(server.running) + + def test_a_model_that_is_still_loading_is_not_ready_yet(self): + # llama binds its port first and answers /health with 503 until the + # model is in memory, so the open port on its own is not the signal. + server = self.server(program=ggml.LLAMA, extra=["--healthy-after", "0.4"]) + started = time.monotonic() + server.serve() + self.assertGreaterEqual(time.monotonic() - started, 0.4) + + def test_a_start_that_never_becomes_ready_gives_up(self): + self.patch_attr(ggml, "STARTUP_TIMEOUT", 0.5) + server = self.server(program=ggml.LLAMA, extra=["--healthy-after", "30"]) + with self.assertRaises(ggml.LocalError): + server.serve() + + def test_stopping_leaves_nothing_running(self): + server = self.server() + server.serve() + server.stop() + self.assertFalse(server.running) + self.assertEqual(server.base_url(), "") + + def test_the_last_thing_it_printed_is_available(self): + server = self.server() + server.serve() + self.assertIn("listening", server.error()) + + def test_asking_what_is_running_does_not_wait_for_a_start(self): + """A model being loaded must not freeze the settings window. + + The interface asks a running server what it is doing while a start is in + flight, and a lock held across the whole start would stop it dead. + """ + server = self.server(extra=["--wait", "0.6"]) + answers = [] + + def start(): + server.serve() + + thread = __import__("threading").Thread(target=start) + thread.start() + try: + time.sleep(0.15) + began = time.monotonic() + answers.append(server.settings()) + answers.append(server.running) + self.assertLess(time.monotonic() - began, 0.2) + finally: + thread.join(timeout=10) + + @linux_only + def test_a_server_a_killed_dikte_left_behind_is_swept_up(self): + server = self.server() + server.serve() + # What a SIGKILL of Dikte leaves: the child still running, the pid file + # still on disk, and nothing left that knows about either. + proc, server._proc = server._proc, None + self.assertTrue(server.sweep()) + self.assertEqual(proc.wait(timeout=5), -signal.SIGTERM) + + @linux_only + def test_a_pid_that_belongs_to_something_else_is_left_alone(self): + server = self.server() + server._remember(os.getpid()) # this test runner, not a server + self.assertFalse(server.sweep()) + + def test_no_pid_file_is_nothing_to_sweep(self): + self.assertFalse(self.server().sweep()) + + +class Arguments(Local): + """What the two command lines say, since neither program is here to say it.""" + + def setUp(self): + super().setUp() + self.binary = self.path("whisper-server") + self.binary.write_text("#!/bin/sh\n") + self.binary.chmod(0o755) + + def whisper_model(self, name="ggml-base.bin"): + path = ggml.whisper_model_path(name) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"model") + return name + + def test_the_inference_path_is_the_one_api_py_builds(self): + args = ggml._whisper_args({"binary": str(self.binary), "gpu": True, + "threads": 0, "model": self.whisper_model()}) + self.assertIn("--inference-path", args) + self.assertEqual(args[args.index("--inference-path") + 1], + "/v1/audio/transcriptions") + + def test_detection_rather_than_english_when_nothing_is_asked_for(self): + args = ggml._whisper_args({"binary": str(self.binary), "gpu": True, + "threads": 0, "model": self.whisper_model()}) + self.assertEqual(args[args.index("-l") + 1], "auto") + + def test_the_graphics_card_is_turned_off_rather_than_asked_for(self): + settings = {"binary": str(self.binary), "gpu": False, "threads": 2, + "model": self.whisper_model()} + args = ggml._whisper_args(settings) + self.assertIn("-ng", args) + self.assertEqual(args[args.index("-t") + 1], "2") + + def test_a_missing_model_is_a_message_about_settings(self): + with self.assertRaises(ggml.LocalError) as caught: + ggml._whisper_args({"binary": str(self.binary), "gpu": True, + "threads": 0, "model": "ggml-nothing.bin"}) + self.assertIn("Settings", str(caught.exception)) + + def test_a_missing_program_says_so_before_a_missing_model(self): + with mock.patch("shutil.which", return_value=None): + with self.assertRaises(ggml.LocalError) as caught: + ggml._whisper_args({"binary": "", "gpu": True, "threads": 0, + "model": self.whisper_model()}) + self.assertIn("whisper.cpp", str(caught.exception)) + + def test_the_layers_go_to_the_card_when_there_is_one(self): + model = ggml.llm_model_path("m.gguf") + model.parent.mkdir(parents=True, exist_ok=True) + model.write_bytes(b"gguf") + args = ggml._llm_args({"binary": str(self.binary), "gpu": True, + "threads": 0, "model": "m.gguf", "context": 4096}) + self.assertEqual(args[args.index("-ngl") + 1], "99") + self.assertEqual(args[args.index("-c") + 1], "4096") + + def test_no_card_means_no_layers_offloaded(self): + model = ggml.llm_model_path("m.gguf") + model.parent.mkdir(parents=True, exist_ok=True) + model.write_bytes(b"gguf") + args = ggml._llm_args({"binary": str(self.binary), "gpu": False, + "threads": 0, "model": "m.gguf", "context": 4096}) + self.assertEqual(args[args.index("-ngl") + 1], "0") + + +class Sizes(DikteTest): + def test_bytes_are_written_the_way_a_download_is_talked_about(self): + self.assertEqual(ggml.human_size(512), "512 B") + self.assertEqual(ggml.human_size(574041195), "547.4 MB") + self.assertEqual(ggml.human_size(3_095_033_483), "2.9 GB") diff --git a/tests/test_hub.py b/tests/test_hub.py new file mode 100644 index 0000000..f3741c2 --- /dev/null +++ b/tests/test_hub.py @@ -0,0 +1,184 @@ +"""What GitHub and Hugging Face are asked, and what is believed of the answer.""" + +import json + +import hub +from tests.support import DikteTest, fake_urlopen, http_error, url_error + +RELEASE = { + "tag_name": "v1.9.1", + "assets": [ + {"name": "whisper-bin-ubuntu-x64.tar.gz", + "browser_download_url": "https://example.invalid/ubuntu-x64.tar.gz", + "size": 9379235, "digest": "sha256:" + "a" * 64}, + {"name": "whisper-bin-x64.zip", + "browser_download_url": "https://example.invalid/win.zip", + "size": 100, "digest": "sha256:" + "b" * 64}, + {"name": "no-url-here.zip", "size": 1}, + ], +} + +TREE = [ + {"type": "file", "path": ".gitattributes", "size": 1477}, + {"type": "file", "path": "ggml-base.bin", "size": 147951465, + "lfs": {"oid": "c" * 64, "size": 147951465}}, + {"type": "directory", "path": "extra"}, + {"type": "file", "path": "extra/ggml-tiny.bin", "size": 77691713, + "lfs": {"oid": "d" * 64, "size": 77691713}}, +] + +MODELS = [ + {"id": "ggml-org/gemma-3-4b-it-GGUF", "downloads": 44606, + "lastModified": "2026-07-01T00:00:00.000Z"}, + {"id": "ggml-org/gpt-oss-20b-GGUF", "downloads": 47975}, + {"noid": True}, +] + + +class Releases(DikteTest): + def setUp(self): + super().setUp() + self.patch_attr(hub, "CACHE_DIR", self.path("cache")) + + def test_the_tag_and_the_assets_come_back(self): + with fake_urlopen(RELEASE) as calls: + tag, assets = hub.release("ggml-org/whisper.cpp") + self.assertEqual(tag, "v1.9.1") + self.assertEqual([a.name for a in assets], + ["whisper-bin-ubuntu-x64.tar.gz", "whisper-bin-x64.zip"]) + self.assertEqual(calls[0].full_url, + "https://api.github.com/repos/ggml-org/whisper.cpp/" + "releases/latest") + + def test_the_sha256_prefix_is_dropped(self): + with fake_urlopen(RELEASE): + _, assets = hub.release("ggml-org/whisper.cpp") + self.assertEqual(assets[0].sha256, "a" * 64) + + def test_a_tag_asks_for_that_tag(self): + with fake_urlopen(RELEASE) as calls: + hub.release("ggml-org/whisper.cpp", "v1.9.1") + self.assertTrue(calls[0].full_url.endswith("/releases/tags/v1.9.1")) + + def test_a_release_with_no_assets_is_an_error(self): + with fake_urlopen({"tag_name": "v1", "assets": []}): + with self.assertRaises(hub.HubError): + hub.release("ggml-org/whisper.cpp") + + def test_the_second_call_asks_nobody(self): + with fake_urlopen(RELEASE) as calls: + hub.release("ggml-org/whisper.cpp") + hub.release("ggml-org/whisper.cpp") + self.assertEqual(len(calls), 1) + + def test_a_refresh_asks_again(self): + with fake_urlopen(RELEASE) as calls: + hub.release("ggml-org/whisper.cpp") + hub.release("ggml-org/whisper.cpp", refresh=True) + self.assertEqual(len(calls), 2) + + def test_an_old_cache_beats_no_answer(self): + with fake_urlopen(RELEASE): + hub.release("ggml-org/whisper.cpp") + # Old enough that it would normally be fetched again, and no network + # to fetch it with. + for path in self.path("cache").iterdir(): + os_utime(path) + with fake_urlopen(url_error()): + tag, assets = hub.release("ggml-org/whisper.cpp") + self.assertEqual(tag, "v1.9.1") + self.assertEqual(len(assets), 2) + + def test_no_cache_and_no_network_says_so(self): + with fake_urlopen(url_error("no route to host")): + with self.assertRaises(hub.HubError) as caught: + hub.release("ggml-org/whisper.cpp") + self.assertIn("api.github.com", str(caught.exception)) + + def test_an_http_error_names_the_host_and_the_code(self): + with fake_urlopen(http_error(404, "nope")): + with self.assertRaises(hub.HubError) as caught: + hub.release("ggml-org/nothing") + self.assertIn("404", str(caught.exception)) + + +class Files(DikteTest): + def setUp(self): + super().setUp() + self.patch_attr(hub, "CACHE_DIR", self.path("cache")) + + def test_directories_are_left_out(self): + with fake_urlopen(TREE): + files = hub.files("ggerganov/whisper.cpp") + self.assertEqual([f.name for f in files], + [".gitattributes", "ggml-base.bin", "extra/ggml-tiny.bin"]) + + def test_the_url_is_the_one_that_serves_the_bytes(self): + with fake_urlopen(TREE): + files = hub.files("ggerganov/whisper.cpp") + self.assertEqual( + files[1].url, + "https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.bin") + + def test_the_lfs_object_id_is_the_checksum(self): + with fake_urlopen(TREE): + files = hub.files("ggerganov/whisper.cpp") + self.assertEqual(files[1].sha256, "c" * 64) + self.assertEqual(files[1].size, 147951465) + + def test_a_file_outside_lfs_has_no_checksum(self): + with fake_urlopen(TREE): + files = hub.files("ggerganov/whisper.cpp") + self.assertEqual(files[0].sha256, "") + + def test_an_answer_that_is_not_a_list_is_an_error(self): + with fake_urlopen({"error": "Invalid username or password."}): + with self.assertRaises(hub.HubError): + hub.files("ggml-org/whisper.cpp") + + +class Repos(DikteTest): + def setUp(self): + super().setUp() + self.patch_attr(hub, "CACHE_DIR", self.path("cache")) + + def test_it_asks_for_one_author_and_for_gguf(self): + with fake_urlopen(MODELS) as calls: + found = hub.repos(author="ggml-org") + self.assertIn("author=ggml-org", calls[0].full_url) + self.assertIn("filter=gguf", calls[0].full_url) + self.assertEqual([r.id for r in found], + ["ggml-org/gemma-3-4b-it-GGUF", "ggml-org/gpt-oss-20b-GGUF"]) + + def test_a_missing_download_count_is_zero(self): + with fake_urlopen(MODELS): + found = hub.repos(author="ggml-org") + self.assertEqual(found[0].downloads, 44606) + self.assertEqual(found[1].updated, "") + + +def os_utime(path): + """Backdate a cache file past its time to live.""" + import os + import time + old = time.time() - hub.CACHE_TTL - 60 + os.utime(path, (old, old)) + + +class CacheOnDisk(DikteTest): + def setUp(self): + super().setUp() + self.patch_attr(hub, "CACHE_DIR", self.path("cache")) + + def test_what_is_stored_is_what_came_back(self): + with fake_urlopen(RELEASE): + hub.release("ggml-org/whisper.cpp") + stored = [json.loads(p.read_text()) for p in self.path("cache").iterdir()] + self.assertEqual(stored[0]["tag_name"], "v1.9.1") + + def test_a_cache_that_cannot_be_written_is_not_a_failure(self): + self.patch_attr(hub, "CACHE_DIR", self.path("nope", "deeper")) + self.path("nope").write_text("a file where a directory would go") + with fake_urlopen(RELEASE): + tag, _ = hub.release("ggml-org/whisper.cpp") + self.assertEqual(tag, "v1.9.1")