wfloat 2.0.0__py3-none-win_amd64.whl

wfloat/_download.py

@@ -0,0 +1,130 @@
+import hashlib
+import os
+import shutil
+import tarfile
+import uuid
+import zipfile
+from pathlib import Path
+from typing import Optional
+from urllib.request import Request, urlopen
+
+
+def normalize_checksum(checksum: str) -> str:
+    normalized = checksum.strip().lower()
+    if normalized.startswith("sha256:"):
+        return normalized.split(":", 1)[1]
+    if normalized.startswith("sha256-"):
+        return normalized.split("-", 1)[1]
+    return normalized
+
+
+def sha256_file(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as file_obj:
+        while True:
+            chunk = file_obj.read(1024 * 1024)
+            if not chunk:
+                break
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+def verify_checksum(path: Path, checksum: str) -> bool:
+    if not path.is_file():
+        return False
+    return sha256_file(path) == normalize_checksum(checksum)
+
+
+def download_file(
+    url: str,
+    destination: Path,
+    *,
+    expected_checksum: Optional[str] = None,
+    timeout: float = 60.0,
+) -> Path:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    temp_path = destination.with_name(destination.name + ".tmp-" + uuid.uuid4().hex)
+
+    request = Request(
+        url,
+        headers={
+            "Accept": "*/*",
+            "User-Agent": "wfloat-python/0.0.0",
+        },
+        method="GET",
+    )
+
+    try:
+        with urlopen(request, timeout=timeout) as response:
+            with temp_path.open("wb") as file_obj:
+                shutil.copyfileobj(response, file_obj)
+
+        if expected_checksum is not None and not verify_checksum(temp_path, expected_checksum):
+            raise RuntimeError(
+                "Downloaded file checksum did not match expected value for %s." % url
+            )
+
+        os.replace(str(temp_path), str(destination))
+        return destination
+    finally:
+        if temp_path.exists():
+            temp_path.unlink()
+
+
+def _assert_within_destination(path: Path, destination: Path) -> None:
+    resolved_destination = destination.resolve()
+    resolved_path = path.resolve()
+    destination_str = str(resolved_destination)
+    path_str = str(resolved_path)
+    if path_str != destination_str and not path_str.startswith(destination_str + os.sep):
+        raise RuntimeError("Archive member would extract outside destination: %s" % path)
+
+
+def _extract_zip(archive_path: Path, destination: Path) -> None:
+    with zipfile.ZipFile(archive_path) as archive:
+        for member in archive.infolist():
+            member_path = destination / member.filename
+            _assert_within_destination(member_path, destination)
+        archive.extractall(destination)
+
+
+def _extract_tar(archive_path: Path, destination: Path) -> None:
+    with tarfile.open(archive_path) as archive:
+        for member in archive.getmembers():
+            member_path = destination / member.name
+            _assert_within_destination(member_path, destination)
+        archive.extractall(destination)
+
+
+def extract_archive(archive_path: Path, destination: Path) -> None:
+    destination.mkdir(parents=True, exist_ok=True)
+    if zipfile.is_zipfile(archive_path):
+        _extract_zip(archive_path, destination)
+        return
+
+    if tarfile.is_tarfile(archive_path):
+        _extract_tar(archive_path, destination)
+        return
+
+    raise RuntimeError(
+        "Unsupported archive format for %s. Python model assets should provide a zip or tar archive."
+        % archive_path
+    )
+
+
+def resolve_extracted_data_directory(extraction_root: Path) -> Path:
+    visible_contents = [
+        path
+        for path in extraction_root.iterdir()
+        if not path.name.startswith(".") and path.name != "__MACOSX"
+    ]
+
+    named_directory = extraction_root / "espeak-ng-data"
+    if named_directory.is_dir():
+        return named_directory
+
+    child_directories = [path for path in visible_contents if path.is_dir()]
+    if len(visible_contents) == 1 and len(child_directories) == 1:
+        return child_directories[0]
+
+    return extraction_root

wfloat/_generated_model_urls.py

@@ -0,0 +1,46 @@
+# Generated from wfloat/assets/registry.json. Do not edit.
+
+REGISTRY_ORIGIN = 'https://registry.wfloat.com'
+
+MODEL_ASSETS = {'HuggingFaceTB/SmolLM2-360M-Instruct': {'family': 'smollm',
+                                         'model': {'path': '/models/huggingfacetb/smollm2-360m-instruct/model.Q4_K_M.gguf',
+                                                   'sha256': '75c4346ef9e855ed630f80078a2430cf63aaca599e340360998a313070fcdc47'}},
+ 'UsefulSensors/moonshine-tiny': {'cached_decoder': {'path': '/models/usefulsensors/moonshine-tiny/cached_decoder.int8.onnx',
+                                                     'sha256': '2aff28bba6a03d8dcf5c9feac45462629bae37317442299f28115ad09da773f6'},
+                                  'encoder': {'path': '/models/usefulsensors/moonshine-tiny/encoder.int8.onnx',
+                                              'sha256': '8774dfba578de027ec6595c2c654a0836434489bc963a0db124a7f181f571acb'},
+                                  'family': 'moonshine',
+                                  'preprocessor': {'path': '/models/usefulsensors/moonshine-tiny/preprocessor.onnx',
+                                                   'sha256': 'f33addce61a143460fe753b5ee5b7db255e5140b5b779c065b94f6c83ff0bf4e'},
+                                  'tokens': {'path': '/models/usefulsensors/moonshine-tiny/tokens.txt',
+                                             'sha256': '1165c2aeb9f72f457a83be2d459a09054f27490acd9b41bd43794dfd25e296ea'},
+                                  'uncached_decoder': {'path': '/models/usefulsensors/moonshine-tiny/uncached_decoder.int8.onnx',
+                                                       'sha256': '216737000dd5881a17aa043f6bbd286add33e4c3b0ae257153e2ec15438bdc41'}},
+ 'k2-fsa/streaming-zipformer-en': {'decoder': {'path': '/models/k2-fsa/streaming-zipformer-en/decoder.onnx',
+                                               'sha256': '9da02b77cb08826756ec6a88635f35a40374e4164e7c6359121a9145958a6ceb'},
+                                   'encoder': {'path': '/models/k2-fsa/streaming-zipformer-en/encoder.int8.onnx',
+                                               'sha256': '32c98281c7bd8b63e3e142d007251b37f120572e8fdea9a4f5a79ce22b10ec4f'},
+                                   'family': 'zipformer-transducer',
+                                   'joiner': {'path': '/models/k2-fsa/streaming-zipformer-en/joiner.onnx',
+                                              'sha256': 'bd5c26ad6a41cbd90c2cfa239c0b55b145af878ce1d79b4739d90f8be93359ba'},
+                                   'tokens': {'path': '/models/k2-fsa/streaming-zipformer-en/tokens.txt',
+                                              'sha256': '49e3c2646595fd907228b3c6787069658f67b17377c60aeb8619c4551b2316fb'}},
+ 'openai/whisper-tiny-en': {'decoder': {'path': '/models/openai/whisper-tiny-en/decoder.int8.onnx',
+                                        'sha256': '06c0e6ff6348d427e51839219d1c886c18cfdf411e629e33f5e1679bff9c1527'},
+                            'encoder': {'path': '/models/openai/whisper-tiny-en/encoder.int8.onnx',
+                                        'sha256': '0ce578b827c94a961aacb8fa14b02f096504b337e5c94be37c36238cbe3e8bc6'},
+                            'family': 'whisper',
+                            'tokens': {'path': '/models/openai/whisper-tiny-en/tokens.txt',
+                                       'sha256': '306cd27f03c1a714eca7108e03d66b7dc042abe8c258b44c199a7ed9838dd930'}},
+ 'snakers4/silero-vad': {'family': 'silero-vad',
+                         'model': {'path': '/models/snakers4/silero-vad/model.onnx',
+                                   'sha256': '9e2449e1087496d8d4caba907f23e0bd3f78d91fa552479bb9c23ac09cbb1fd6'}},
+ 'wfloat/wfloat-tts': {'model_onnx': {'path': '/models/wfloat/wfloat-tts/model.onnx',
+                                      'sha256': 'a7e65773a29499b80a393bbe08af3507e18f6ef95faa0eaf7cb4ba353c8693ae'},
+                       'model_tokens': {'path': '/models/wfloat/wfloat-tts/tokens.txt',
+                                        'sha256': '96fd291bede0544469d4d8935d462fdd6dc947f22ad47369753e1a82db3d748e'}}}
+
+SHARED_ASSETS = {'espeak_ng_data_aar': {'path': '/assets/espeak-ng-data/espeak-ng-data-2023.9.7-4.aar',
+                        'sha256': 'a526b72e81cb1a17e07f55ca0117bba8fbcac7ccd2fa502c61be926eafeaf64e'},
+ 'espeak_ng_data_zip': {'path': '/assets/espeak-ng-data/espeak-ng-data-2023.9.7-4.zip',
+                        'sha256': '56c2879ab1ab44c594c78f34e76c50cf1dd7b8f6ca0ca2634b6766a6edb32add'}}

wfloat/_llm.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Dict, Optional, Sequence
+
+from ._results import LlmGenerationResult
+
+
+@dataclass
+class LlmModel:
+    model_id: str
+    family: str
+    _native_llm: object
+    context_size: int = 0
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        max_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repeat_penalty: float = 1.0,
+        seed: int = 0,
+        on_token: Optional[Callable[[str], None]] = None,
+    ) -> LlmGenerationResult:
+        if not hasattr(self._native_llm, "generate"):
+            raise RuntimeError("Native LLM backend does not support generate().")
+
+        return self._native_llm.generate(
+            prompt,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repeat_penalty=repeat_penalty,
+            seed=seed,
+            on_token=on_token,
+        )
+
+    def chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        *,
+        max_tokens: int = 128,
+        temperature: float = 0.8,
+        top_p: float = 0.95,
+        top_k: int = 40,
+        repeat_penalty: float = 1.0,
+        seed: int = 0,
+        on_token: Optional[Callable[[str], None]] = None,
+    ) -> LlmGenerationResult:
+        if not hasattr(self._native_llm, "chat"):
+            raise RuntimeError("Native LLM backend does not support chat().")
+
+        return self._native_llm.chat(
+            messages,
+            max_tokens=max_tokens,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            repeat_penalty=repeat_penalty,
+            seed=seed,
+            on_token=on_token,
+        )
+
+    def format_chat(
+        self,
+        messages: Sequence[Dict[str, str]],
+        *,
+        add_generation_prompt: bool = True,
+    ) -> str:
+        if not hasattr(self._native_llm, "format_chat"):
+            raise RuntimeError("Native LLM backend does not support format_chat().")
+
+        return self._native_llm.format_chat(
+            messages,
+            add_generation_prompt=add_generation_prompt,
+        )
+
+    def close(self) -> None:
+        if hasattr(self._native_llm, "close"):
+            self._native_llm.close()

wfloat/_llm_assets.py

@@ -0,0 +1,142 @@
+from __future__ import annotations
+
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Mapping, Optional
+from urllib.parse import urlparse
+
+from ._assets import LlmModelAssets
+from ._download import download_file, verify_checksum
+
+
+@dataclass(frozen=True)
+class CachedLlmAssets:
+    model_name: str
+    family: str
+    cache_dir: Path
+    files: Mapping[str, Path]
+    context_size: Optional[int] = None
+    chat_template: Optional[str] = None
+    chat_template_format: Optional[str] = None
+
+    def require(self, key: str) -> Path:
+        path = self.files.get(key)
+        if path is None:
+            raise ValueError(f"Missing required LLM asset: {key}")
+        return path
+
+
+def _normalize_model_dir_name(model_name: str) -> str:
+    return model_name.replace("/", "--").replace(" ", "-")
+
+
+def _is_url(value: str) -> bool:
+    parsed = urlparse(value)
+    return parsed.scheme in {"http", "https", "file"}
+
+
+def _copy_local_file(source: Path, destination: Path) -> None:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copyfile(str(source), str(destination))
+
+
+def _materialize_asset(
+    source: str | Path,
+    destination: Path,
+    *,
+    expected_checksum: Optional[str],
+    force_download: bool,
+) -> Path:
+    destination.parent.mkdir(parents=True, exist_ok=True)
+
+    if not force_download and destination.is_file():
+        if expected_checksum is None or verify_checksum(destination, expected_checksum):
+            return destination
+
+    if isinstance(source, Path):
+        _copy_local_file(source, destination)
+    else:
+        source_str = str(source)
+        if _is_url(source_str):
+            download_file(source_str, destination, expected_checksum=expected_checksum)
+        else:
+            _copy_local_file(Path(source_str), destination)
+
+    if expected_checksum is not None and not verify_checksum(destination, expected_checksum):
+        raise RuntimeError(f"Cached LLM asset checksum mismatch for {destination}.")
+
+    return destination
+
+
+def cache_llm_assets(
+    model_name: str,
+    *,
+    family: str,
+    sources: Mapping[str, str | Path | None],
+    checksums: Optional[Mapping[str, str]] = None,
+    cache_dir: Path,
+    force_download: bool = False,
+    context_size: Optional[int] = None,
+    chat_template: Optional[str] = None,
+    chat_template_format: Optional[str] = None,
+) -> CachedLlmAssets:
+    model_dir = cache_dir / "models" / _normalize_model_dir_name(model_name)
+    model_dir.mkdir(parents=True, exist_ok=True)
+    checksums = checksums or {}
+
+    files: dict[str, Path] = {}
+    for key, source in sources.items():
+        if source is None:
+            continue
+
+        source_value = Path(source) if isinstance(source, Path) else str(source)
+        filename = Path(urlparse(str(source_value)).path).name or Path(str(source_value)).name
+        if not filename:
+            raise ValueError(f"Could not derive filename for LLM asset {key}.")
+
+        destination = model_dir / filename
+        files[key] = _materialize_asset(
+            source_value,
+            destination,
+            expected_checksum=checksums.get(key),
+            force_download=force_download,
+        )
+
+    return CachedLlmAssets(
+        model_name=model_name,
+        family=family,
+        cache_dir=cache_dir,
+        files=files,
+        context_size=context_size,
+        chat_template=chat_template,
+        chat_template_format=chat_template_format,
+    )
+
+
+def cache_llm_model_assets(
+    model_name: str,
+    assets: LlmModelAssets,
+    *,
+    cache_dir: Path,
+    force_download: bool = False,
+) -> CachedLlmAssets:
+    return cache_llm_assets(
+        model_name,
+        family=assets.family,
+        sources={
+            "model": assets.model,
+        },
+        checksums={
+            key: value
+            for key, value in {
+                "model": assets.model_checksum,
+            }.items()
+            if value is not None
+        },
+        cache_dir=cache_dir,
+        force_download=force_download,
+        context_size=assets.context_size,
+        chat_template=assets.chat_template,
+        chat_template_format=assets.chat_template_format,
+    )

wfloat/_llm_load.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+from ._assets import fetch_llm_assets
+from ._cache import get_default_cache_dir
+from ._core import create_core_llm
+from ._llm import LlmModel
+from ._llm_assets import cache_llm_model_assets
+
+DEFAULT_LLM_CONTEXT_SIZE = 2048
+DEFAULT_LLM_NUM_THREADS = 4
+DEFAULT_LLM_GPU_LAYER_COUNT = 0
+
+
+def load_llm_model(
+    model_name: str,
+    *,
+    cache_dir: Optional[Path] = None,
+    force_download: bool = False,
+    context_size: Optional[int] = None,
+    num_threads: int = DEFAULT_LLM_NUM_THREADS,
+    gpu_layer_count: int = DEFAULT_LLM_GPU_LAYER_COUNT,
+    chat_template: Optional[str] = None,
+) -> LlmModel:
+    resolved_cache_dir = Path(cache_dir) if cache_dir is not None else get_default_cache_dir()
+    assets = fetch_llm_assets(model_name)
+    cached = cache_llm_model_assets(
+        model_name,
+        assets,
+        cache_dir=resolved_cache_dir,
+        force_download=force_download,
+    )
+    family = assets.family
+
+    resolved_context_size = (
+        int(context_size)
+        if context_size is not None
+        else int(cached.context_size or DEFAULT_LLM_CONTEXT_SIZE)
+    )
+    resolved_chat_template = chat_template if chat_template is not None else cached.chat_template
+    if resolved_chat_template is None and getattr(cached, "chat_template_format", None) == "chatml":
+        resolved_chat_template = "chatml"
+
+    native_llm = create_core_llm(
+        model_name=model_name,
+        family=family,
+        model_path=cached.require("model"),
+        context_size=resolved_context_size,
+        num_threads=num_threads,
+        gpu_layer_count=gpu_layer_count,
+        chat_template=resolved_chat_template,
+    )
+
+    return LlmModel(
+        model_id=model_name,
+        family=family,
+        _native_llm=native_llm,
+        context_size=resolved_context_size,
+    )