funasr-mlx 0.1.0__py3-none-any.whl

funasr_mlx/__init__.py

@@ -0,0 +1,3 @@
+"""Port notes: docs/port-notes/__init__.py.md"""
+
+__version__ = "0.1.0"

funasr_mlx/api.py

@@ -0,0 +1,92 @@
+"""Port notes: docs/port-notes/api.py.md"""
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+from pathlib import Path
+from typing import Any
+
+from .memory import MemoryBudgetExceeded
+
+def download_model(repo, *, variant=None, dest=".", owner=None):
+    # see docs/port-notes/functions.md :: download_model
+    from .download import download
+
+    kwargs = {"variant": variant, "dest": dest}
+    if owner is not None:
+        kwargs["owner"] = owner
+    return download(repo, **kwargs)
+
+
+def pseudo_stream(model, *, family: str | None = None, chunk_seconds: float = 1.0):
+    # see docs/port-notes/functions.md :: pseudo_stream
+    from .pseudo_streaming import PseudoStream
+
+    transcriber = model if hasattr(model, "transcribe_samples") else TranscribeWorker(
+        model, family=family
+    )._impl
+    return PseudoStream(
+        transcribe_fn=transcriber.transcribe_samples,
+        chunk_seconds=chunk_seconds,
+    )
+
+
+__all__ = ["transcribe", "TranscribeWorker", "download_model", "pseudo_stream"]
+
+
+class TranscribeWorker:
+
+    def __init__(self, checkpoint: str | Path, *, family: str | None = None, max_tokens: int = 96):
+        from .llm_asr import LlmAsrTranscriber, resolve_family
+
+        checkpoint = Path(checkpoint)
+        self._family = resolve_family(checkpoint, family)
+        self._impl = LlmAsrTranscriber(checkpoint, family=self._family, max_tokens=max_tokens)
+
+    @property
+    def family(self) -> str:
+        return self._family
+
+    def transcribe(self, audio: str | Path) -> dict[str, Any]:
+        return self._impl.transcribe_file(audio).to_dict()
+
+    def pseudo_stream(self, chunk_seconds: float = 1.0):
+        from .pseudo_streaming import PseudoStream
+
+        return PseudoStream(
+            transcribe_fn=self._impl.transcribe_samples,
+            chunk_seconds=chunk_seconds,
+        )
+
+
+def transcribe(
+    audio: str | Path,
+    *,
+    model: str | Path,
+    family: str | None = None,
+    max_tokens: int = 96,
+) -> dict[str, Any]:
+    try:
+        worker = TranscribeWorker(model, family=family, max_tokens=max_tokens)
+        return worker.transcribe(audio)
+    except MemoryBudgetExceeded:
+        pass  # fall through to the guarded CLI subprocess
+
+    cmd = [sys.executable, "-m", "funasr_mlx.cli", "transcribe", str(model),
+           "--audio", str(audio), "--max-tokens", str(max_tokens)]
+    if family:
+        cmd += ["--family", family]
+    env = dict(os.environ)
+    env.pop("FUNASR_MLX_MEMORY_POLICY", None)  # let the CLI parent set its own
+    env["PYTHONPATH"] = str(Path(__file__).resolve().parent.parent)
+    proc = subprocess.Popen(
+        cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, env=env
+    )
+    out, err = proc.communicate(timeout=900)
+    for line in out.splitlines():
+        line = line.strip()
+        if line.startswith("{") and '"text"' in line:
+            return json.loads(line)
+    raise RuntimeError(f"guarded transcribe failed (rc={proc.returncode}): {err[-400:]}")

funasr_mlx/checkpoint_stream.py

@@ -0,0 +1,138 @@
+"""Port notes: docs/port-notes/checkpoint_stream.py.md"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+import struct
+
+import numpy as np
+
+_CODES = {
+    "float32": "F32",
+    "float16": "F16",
+    "float64": "F64",
+    "int64": "I64",
+    "int32": "I32",
+    "int16": "I16",
+    "int8": "I8",
+    "uint8": "U8",
+    "uint32": "U32",
+    "uint64": "U64",
+    "bool": "BOOL",
+}
+
+# see docs/port-notes/comments.md :: checkpoint_stream.py:25
+_CODE_BYTES = {
+    "BF16": 2,
+    "F16": 2,
+    "F32": 4,
+    "F64": 8,
+    "I8": 1,
+    "U8": 1,
+    "I16": 2,
+    "I32": 4,
+    "I64": 8,
+    "U32": 4,
+    "U64": 8,
+    "BOOL": 1,
+}
+
+
+def write_shard(
+    path: Path,
+    tensors: dict,
+    *,
+    metadata: dict,
+    code: str | None = None,
+    codes: dict | None = None,
+) -> None:
+    """Write a whole shard whose tensors are all carried in one dtype.
+
+    Full notes: docs/port-notes/functions.md :: write_shard
+    """
+    if (code is None) == (codes is None):
+        raise ValueError("Pass exactly one of code or codes")
+    if code is not None:
+        if code not in _CODE_BYTES:
+            raise ValueError(f"Unsupported shard type code {code}")
+        codes = {name: code for name in tensors}
+    header: dict = {"__metadata__": dict(metadata)}
+    offset = 0
+    for name, array in tensors.items():
+        tensor_code = codes.get(name)
+        if tensor_code not in _CODE_BYTES:
+            raise ValueError(f"Unsupported shard type code {tensor_code} for {name}")
+        width = _CODE_BYTES[tensor_code]
+        if array.dtype.itemsize != width:
+            raise ValueError(
+                f"Shard tensor {name} has element width {array.dtype.itemsize}, "
+                f"but code {tensor_code} needs {width}"
+            )
+        size = int(array.size) * width
+        header[name] = {
+            "dtype": tensor_code,
+            "shape": list(array.shape),
+            "data_offsets": [offset, offset + size],
+        }
+        offset += size
+    encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    encoded += b" " * ((-len(encoded)) % 8)
+    temporary = path.with_suffix(path.suffix + ".partial")
+    try:
+        with temporary.open("wb") as output:
+            output.write(struct.pack("<Q", len(encoded)))
+            output.write(encoded)
+            for array in tensors.values():
+                output.write(np.ascontiguousarray(array).tobytes())
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)
+
+
+def write_tensor(
+    path: Path,
+    name: str,
+    array: np.ndarray,
+    *,
+    dtype,
+    metadata,
+    chunk_bytes: int = 16_000_000,
+    code: str | None = None,
+):
+    # see docs/port-notes/functions.md :: write_tensor
+    output_dtype = np.dtype(dtype)
+    if not array.flags.c_contiguous:
+        raise ValueError("Large streaming tensors must be C-contiguous")
+    if code is not None:
+        if _CODE_BYTES.get(code) != array.dtype.itemsize:
+            raise ValueError("Streaming type code disagrees with the array element width")
+        output_code = code
+    elif not output_dtype.isnative or output_dtype.name not in _CODES:
+        raise ValueError("Unsupported streaming tensor dtype")
+    else:
+        output_code = _CODES[output_dtype.name]
+    count = int(array.size)
+    header = {
+        "__metadata__": dict(metadata),
+        name: {
+            "dtype": output_code,
+            "shape": list(array.shape),
+            "data_offsets": [0, count * output_dtype.itemsize],
+        },
+    }
+    encoded = json.dumps(header, separators=(",", ":")).encode("utf-8")
+    encoded += b" " * ((-len(encoded)) % 8)
+    temporary = path.with_suffix(path.suffix + ".partial")
+    flat = array.reshape(-1)
+    stride = max(1, chunk_bytes // max(array.dtype.itemsize, output_dtype.itemsize))
+    try:
+        with temporary.open("wb") as output:
+            output.write(struct.pack("<Q", len(encoded)))
+            output.write(encoded)
+            for start in range(0, count, stride):
+                block = np.asarray(flat[start : start + stride], dtype=output_dtype, order="C")
+                output.write(memoryview(block).cast("B"))
+        temporary.replace(path)
+    finally:
+        temporary.unlink(missing_ok=True)