modeling.py

from __future__ import annotations

import json
import os
from typing import Optional

import torch

TARGET_SR = 16000


def validate_audio(wav: torch.Tensor, sr: int):
    if sr != TARGET_SR:
        raise ValueError(
            f"sample rate must be {TARGET_SR} Hz after resampling; got {sr}"
        )
    if wav.dim() != 2 or wav.size(0) != 1:
        raise ValueError(
            f"waveform must be mono with shape (1, num_samples); got {tuple(wav.shape)}"
        )
    if wav.numel() == 0:
        raise ValueError("empty waveform")
    if wav.abs().max() > 1.5:
        raise ValueError(
            "waveform looks unnormalized (values outside [-1, 1]); "
            "scale int16 by 1/32768 before passing it in"
        )


def preprocess_audio(source, sr: Optional[int] = None) -> torch.Tensor:
    import torchaudio

    if isinstance(source, str):
        wav, in_sr = torchaudio.load(source)
    else:
        wav = torch.as_tensor(source, dtype=torch.float32)
        if wav.dim() == 1:
            wav = wav.unsqueeze(0)
        in_sr = sr if sr is not None else TARGET_SR
        if wav.abs().max() > 1.5:
            wav = wav / 32768.0
    if wav.size(0) > 1:
        wav = wav.mean(0, keepdim=True)
    if in_sr != TARGET_SR:
        wav = torchaudio.functional.resample(wav, in_sr, TARGET_SR)
    validate_audio(wav, TARGET_SR)
    return wav


def _read_config(config_path: Optional[str]) -> dict:
    if config_path and os.path.exists(config_path):
        with open(config_path) as f:
            return json.load(f)
    return {}


def load_model(
    weights: str, bpe: str = "bpe256.model", config_path: Optional[str] = None
):
    here = os.path.dirname(os.path.abspath(__file__))
    if config_path is None:
        cand = os.path.join(here, "config.json")
        config_path = cand if os.path.exists(cand) else None
    cfg_json = _read_config(config_path)

    from fela_ctc2 import FELACTC2, BPE, greedy_decode_bpe
    from model_cpu_gpt2 import CPUGPTConfig

    bpe_obj = BPE(model_file=bpe)
    cfg = CPUGPTConfig(
        vocab_size=cfg_json.get("vocab", 257),
        seq_len=cfg_json.get("seq_len", 2048),
        n_layer=cfg_json.get("n_layer", 16),
        n_embd=cfg_json.get("n_embd", 512),
        n_head=cfg_json.get("n_head", 8),
        fno_modes=cfg_json.get("fno_modes", 256),
        gla_chunk=cfg_json.get("gla_chunk", 64),
        ffn_hidden=cfg_json.get("ffn_hidden", 2048),
        layer_pattern=cfg_json.get("layer_pattern", "FNO"),
        dropout=0.0,
    )
    if hasattr(cfg, "gla_delta"):
        cfg.gla_delta = bool(cfg_json.get("gla_delta", False))
    model = FELACTC2(cfg, vocab=bpe_obj.vocab).eval()

    if weights.endswith(".safetensors"):
        from safetensors.torch import load_file

        state = load_file(weights)
    else:
        ck = torch.load(weights, map_location="cpu", weights_only=False)
        state = ck.get("state", ck.get("model", ck)) if isinstance(ck, dict) else ck
    model.load_state_dict(state, strict=False)
    return model, bpe_obj, greedy_decode_bpe


def from_pretrained(repo_id: str = "lowdown-labs/asr-streaming"):
    from huggingface_hub import hf_hub_download

    cfg_path = hf_hub_download(repo_id, "config.json")
    bpe = hf_hub_download(repo_id, "bpe256.model")
    try:
        w = hf_hub_download(repo_id, "model.safetensors")
    except Exception:
        w = hf_hub_download(repo_id, "model.pt")
    return load_model(w, bpe=bpe, config_path=cfg_path)


@torch.no_grad()
def stream_transcribe(
    model,
    bpe,
    decode_fn,
    wav: torch.Tensor,
    frame_chunk: int = 100,
    reset_frames: int = 1200,
) -> str:
    out = []
    for logp in model.stream_logits(
        wav, frame_chunk=frame_chunk, reset_frames=reset_frames
    ):
        t = decode_fn(logp[0], bpe)
        if t:
            out.append(t)
    return " ".join(out)