ffi_python/liteasr_ffi.py

import ctypes
from ctypes import POINTER, Structure, c_bool, c_char_p, c_float, c_int32, c_int64, c_size_t, c_uint8, c_uint32
from pathlib import Path

import numpy as np


class LiteAsrError(RuntimeError):
    pass


class _F32Buffer(Structure):
    _fields_ = [
        ("data", POINTER(c_float)),
        ("len", c_size_t),
        ("cap", c_size_t),
        ("rows", c_size_t),
        ("cols", c_size_t),
    ]


class _I64Buffer(Structure):
    _fields_ = [
        ("data", POINTER(c_int64)),
        ("len", c_size_t),
        ("cap", c_size_t),
    ]


class _Utf8Buffer(Structure):
    _fields_ = [
        ("data", POINTER(c_uint8)),
        ("len", c_size_t),
        ("cap", c_size_t),
    ]


def _default_dll_candidates() -> list[Path]:
    return [
        Path("target/debug/liteasr_ffi.dll"),
        Path("target/release/liteasr_ffi.dll"),
    ]


class LiteAsrFfi:
    def __init__(self, dll_path: str | Path | None = None) -> None:
        if dll_path is None:
            for candidate in _default_dll_candidates():
                if candidate.exists():
                    dll_path = candidate
                    break
            if dll_path is None:
                raise LiteAsrError("liteasr_ffi.dll not found. Build with `cargo build`.")

        self._dll_path = Path(dll_path)
        if not self._dll_path.exists():
            raise LiteAsrError(f"DLL not found: {self._dll_path}")

        self._lib = ctypes.CDLL(str(self._dll_path))
        self._configure_signatures()

    @property
    def dll_path(self) -> Path:
        return self._dll_path

    def _configure_signatures(self) -> None:
        self._lib.liteasr_last_error_message.argtypes = []
        self._lib.liteasr_last_error_message.restype = c_char_p

        self._lib.liteasr_preprocess_wav.argtypes = [
            c_char_p,
            c_uint32,
            c_size_t,
            POINTER(_F32Buffer),
        ]
        self._lib.liteasr_preprocess_wav.restype = c_int32

        self._lib.liteasr_build_prompt_ids.argtypes = [
            c_char_p,
            c_char_p,
            c_char_p,
            c_bool,
            c_bool,
            c_bool,
            POINTER(_I64Buffer),
        ]
        self._lib.liteasr_build_prompt_ids.restype = c_int32

        self._lib.liteasr_decode_tokens.argtypes = [
            c_char_p,
            POINTER(c_int64),
            c_size_t,
            c_bool,
            POINTER(_Utf8Buffer),
        ]
        self._lib.liteasr_decode_tokens.restype = c_int32

        self._lib.liteasr_apply_suppression.argtypes = [
            POINTER(c_float),
            c_size_t,
            POINTER(c_int64),
            c_size_t,
            POINTER(c_int64),
            c_size_t,
            c_size_t,
        ]
        self._lib.liteasr_apply_suppression.restype = c_int32

        self._lib.liteasr_free_f32_buffer.argtypes = [POINTER(_F32Buffer)]
        self._lib.liteasr_free_f32_buffer.restype = None
        self._lib.liteasr_free_i64_buffer.argtypes = [POINTER(_I64Buffer)]
        self._lib.liteasr_free_i64_buffer.restype = None
        self._lib.liteasr_free_utf8_buffer.argtypes = [POINTER(_Utf8Buffer)]
        self._lib.liteasr_free_utf8_buffer.restype = None

    def _raise_last_error(self, fallback: str) -> None:
        message = self._lib.liteasr_last_error_message()
        if message:
            text = message.decode("utf-8", errors="replace")
            raise LiteAsrError(text)
        raise LiteAsrError(fallback)

    def preprocess_wav(self, wav_path: str | Path, target_sr: int, n_mels: int) -> np.ndarray:
        out = _F32Buffer()
        rc = self._lib.liteasr_preprocess_wav(
            str(Path(wav_path)).encode("utf-8"),
            int(target_sr),
            int(n_mels),
            ctypes.byref(out),
        )
        if rc != 0:
            self._raise_last_error("liteasr_preprocess_wav failed")

        try:
            flat = np.ctypeslib.as_array(out.data, shape=(out.len,))
            matrix = flat.reshape((out.rows, out.cols)).copy()
            return matrix
        finally:
            self._lib.liteasr_free_f32_buffer(ctypes.byref(out))

    def build_prompt_ids(
        self,
        tokenizer_json_path: str | Path,
        language: str,
        task: str,
        with_timestamps: bool,
        omit_language_token: bool,
        omit_notimestamps_token: bool,
    ) -> list[int]:
        out = _I64Buffer()
        rc = self._lib.liteasr_build_prompt_ids(
            str(Path(tokenizer_json_path)).encode("utf-8"),
            language.encode("utf-8"),
            task.encode("utf-8"),
            bool(with_timestamps),
            bool(omit_language_token),
            bool(omit_notimestamps_token),
            ctypes.byref(out),
        )
        if rc != 0:
            self._raise_last_error("liteasr_build_prompt_ids failed")

        try:
            arr = np.ctypeslib.as_array(out.data, shape=(out.len,))
            return [int(v) for v in arr.tolist()]
        finally:
            self._lib.liteasr_free_i64_buffer(ctypes.byref(out))

    def decode_tokens(
        self,
        tokenizer_json_path: str | Path,
        token_ids: list[int],
        skip_special_tokens: bool = True,
    ) -> str:
        out = _Utf8Buffer()
        token_np = np.array(token_ids, dtype=np.int64)
        token_ptr = token_np.ctypes.data_as(POINTER(c_int64))
        rc = self._lib.liteasr_decode_tokens(
            str(Path(tokenizer_json_path)).encode("utf-8"),
            token_ptr,
            int(token_np.shape[0]),
            bool(skip_special_tokens),
            ctypes.byref(out),
        )
        if rc != 0:
            self._raise_last_error("liteasr_decode_tokens failed")

        try:
            data = ctypes.string_at(out.data, out.len)
            return data.decode("utf-8", errors="replace")
        finally:
            self._lib.liteasr_free_utf8_buffer(ctypes.byref(out))

    def apply_suppression(
        self,
        logits: np.ndarray,
        suppress_ids: list[int],
        begin_suppress_ids: list[int],
        step: int,
    ) -> np.ndarray:
        if logits.dtype != np.float32 or not logits.flags["C_CONTIGUOUS"]:
            logits = np.ascontiguousarray(logits, dtype=np.float32)

        suppress_np = np.array(suppress_ids or [], dtype=np.int64)
        begin_np = np.array(begin_suppress_ids or [], dtype=np.int64)

        suppress_ptr = (
            suppress_np.ctypes.data_as(POINTER(c_int64))
            if suppress_np.size > 0
            else ctypes.cast(0, POINTER(c_int64))
        )
        begin_ptr = (
            begin_np.ctypes.data_as(POINTER(c_int64))
            if begin_np.size > 0
            else ctypes.cast(0, POINTER(c_int64))
        )

        rc = self._lib.liteasr_apply_suppression(
            logits.ctypes.data_as(POINTER(c_float)),
            int(logits.shape[0]),
            suppress_ptr,
            int(suppress_np.shape[0]),
            begin_ptr,
            int(begin_np.shape[0]),
            int(step),
        )
        if rc != 0:
            self._raise_last_error("liteasr_apply_suppression failed")
        return logits