Upload app.py with huggingface_hub

app.py

@@ -1,17 +1,12 @@
-"""
-SozKZ -- Kazakh ASR Demo
-OmniAudio v2 Scratch 70M
-"""
+"""SozKZ -- Kazakh ASR Demo. Uses original model_v2.py from HF repo."""
 
 import os
-import math
 import spaces
 import gradio as gr
 import torch
-import torch.nn as nn
-import torch.nn.functional as F
 import numpy as np
 import librosa
+import soundfile as sf
 import time
 from transformers import PreTrainedTokenizerFast
 from huggingface_hub import hf_hub_download, login
@@ -20,230 +15,68 @@ HF_TOKEN = os.environ.get("HF_TOKEN")
 if HF_TOKEN:
     login(token=HF_TOKEN)
 
-# -- Model (exact names matching model.pt state_dict) --
-
-class RotaryEmbedding(nn.Module):
-    def __init__(self, dim, base=10000.0):
-        super().__init__()
-        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
-        self.register_buffer("inv_freq", inv_freq)
-    def forward(self, seq_len):
-        t = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
-        freqs = torch.outer(t, self.inv_freq)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        return emb.cos(), emb.sin()
-
-def _rotate_half(x):
-    x1, x2 = x[..., :x.shape[-1]//2], x[..., x.shape[-1]//2:]
-    return torch.cat([-x2, x1], dim=-1)
-
-def apply_rotary_emb(x, cos, sin):
-    s = x.shape[2]
-    return x * cos[:s].unsqueeze(0).unsqueeze(0) + _rotate_half(x) * sin[:s].unsqueeze(0).unsqueeze(0)
-
-class RMSNorm(nn.Module):
-    def __init__(self, dim, eps=1e-6):
-        super().__init__()
-        self.eps = eps
-        self.weight = nn.Parameter(torch.ones(dim))
-    def forward(self, x):
-        return (x.float() * x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()).to(x.dtype) * self.weight
-
-class EncoderBlock(nn.Module):
-    def __init__(self, d_model, n_heads, dropout=0.1):
-        super().__init__()
-        self.n_heads = n_heads
-        self.head_dim = d_model // n_heads
-        self.norm1 = RMSNorm(d_model)
-        self.norm2 = RMSNorm(d_model)
-        self.q_proj = nn.Linear(d_model, d_model)
-        self.k_proj = nn.Linear(d_model, d_model)
-        self.v_proj = nn.Linear(d_model, d_model)
-        self.o_proj = nn.Linear(d_model, d_model)
-        self.rope = RotaryEmbedding(self.head_dim)
-        inter = int(d_model * 8 / 3)
-        inter = ((inter + 63) // 64) * 64
-        self.gate_proj = nn.Linear(d_model, inter, bias=False)
-        self.up_proj = nn.Linear(d_model, inter, bias=False)
-        self.down_proj = nn.Linear(inter, d_model, bias=False)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x):
-        B, T, C = x.shape
-        h = self.norm1(x)
-        q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        k = self.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        cos, sin = self.rope(T)
-        q = apply_rotary_emb(q, cos, sin)
-        k = apply_rotary_emb(k, cos, sin)
-        attn = F.scaled_dot_product_attention(q, k, v)
-        x = x + self.dropout(self.o_proj(attn.transpose(1, 2).contiguous().view(B, T, C)))
-        h = self.norm2(x)
-        x = x + self.dropout(self.down_proj(F.silu(self.gate_proj(h)) * self.up_proj(h)))
-        return x
-
-class AudioEncoder(nn.Module):
-    def __init__(self, n_mels=80, d_model=256, n_heads=4, n_layers=6, n_conv=2, dropout=0.1):
-        super().__init__()
-        convs = []
-        inch = n_mels
-        for i in range(n_conv):
-            convs += [nn.Conv1d(inch, d_model, 3, 2, 1), nn.SiLU(), nn.Dropout(dropout)]
-            inch = d_model
-        self.conv_stack = nn.Sequential(*convs)
-        self.layers = nn.ModuleList([EncoderBlock(d_model, n_heads, dropout) for _ in range(n_layers)])
-        self.norm = RMSNorm(d_model)
-
-    def forward(self, mel):
-        x = self.conv_stack(mel).transpose(1, 2)
-        for layer in self.layers:
-            x = layer(x)
-        return self.norm(x)
-
-class DecoderBlock(nn.Module):
-    def __init__(self, d_model, n_heads, dropout=0.1):
-        super().__init__()
-        self.n_heads = n_heads
-        self.head_dim = d_model // n_heads
-        self.norm1 = RMSNorm(d_model)
-        self.norm2 = RMSNorm(d_model)
-        self.q_proj = nn.Linear(d_model, d_model)
-        self.k_proj = nn.Linear(d_model, d_model)
-        self.v_proj = nn.Linear(d_model, d_model)
-        self.o_proj = nn.Linear(d_model, d_model)
-        inter = int(d_model * 8 / 3)
-        inter = ((inter + 63) // 64) * 64
-        self.gate_proj = nn.Linear(d_model, inter, bias=False)
-        self.up_proj = nn.Linear(d_model, inter, bias=False)
-        self.down_proj = nn.Linear(inter, d_model, bias=False)
-        self.dropout = nn.Dropout(dropout)
-
-    def forward(self, x, cos, sin):
-        B, T, C = x.shape
-        h = self.norm1(x)
-        q = self.q_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        k = self.k_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        v = self.v_proj(h).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
-        q = apply_rotary_emb(q, cos, sin)
-        k = apply_rotary_emb(k, cos, sin)
-        attn = F.scaled_dot_product_attention(q, k, v, is_causal=True)
-        x = x + self.dropout(self.o_proj(attn.transpose(1, 2).contiguous().view(B, T, C)))
-        h = self.norm2(x)
-        x = x + self.dropout(self.down_proj(F.silu(self.gate_proj(h)) * self.up_proj(h)))
-        return x
-
-class AudioProjectorV2(nn.Module):
-    def __init__(self, audio_dim, llm_dim):
-        super().__init__()
-        self.linear = nn.Linear(audio_dim, llm_dim)
-        self.norm = RMSNorm(llm_dim)
-    def forward(self, x):
-        return self.norm(self.linear(x))
-
-
-class OmniAudioScratchModel(nn.Module):
-    def __init__(self, encoder_config, decoder_config, vocab_size=50257, dropout=0.1):
-        super().__init__()
-        enc_dim = encoder_config["d_model"]
-        dec_dim = decoder_config["d_model"]
-        self.encoder = AudioEncoder(**encoder_config, dropout=dropout)
-        self.projector = AudioProjectorV2(enc_dim, dec_dim)
-        self.embed_tokens = nn.Embedding(vocab_size, dec_dim)
-        self.decoder_layers = nn.ModuleList([
-            DecoderBlock(dec_dim, decoder_config["n_heads"], dropout)
-            for _ in range(decoder_config["n_layers"])
-        ])
-        self.decoder_norm = RMSNorm(dec_dim)
-        self.decoder_rope = RotaryEmbedding(dec_dim // decoder_config["n_heads"])
-        self.lm_head = nn.Linear(dec_dim, vocab_size, bias=False)
-        # CTC head (may exist in checkpoint, not used for inference)
-        self.ctc_head = nn.Linear(enc_dim, vocab_size)
-
-    def generate(self, mel, max_new_tokens=200, eos_token_id=0, repetition_penalty=1.2):
-        enc_out = self.encoder(mel)
-        audio_embeds = self.projector(enc_out)
-        generated = []
-        combined = audio_embeds
-        for _ in range(max_new_tokens):
-            cos, sin = self.decoder_rope(combined.size(1))
-            x = combined
-            for layer in self.decoder_layers:
-                x = layer(x, cos, sin)
-            logits = self.lm_head(self.decoder_norm(x)[:, -1:]).squeeze(0).squeeze(0)
-            if repetition_penalty != 1.0 and generated:
-                for t in set(generated):
-                    if logits[t] > 0:
-                        logits[t] /= repetition_penalty
-                    else:
-                        logits[t] *= repetition_penalty
-            tok = logits.argmax(-1).item()
-            if tok == eos_token_id:
-                break
-            generated.append(tok)
-            combined = torch.cat([combined, self.embed_tokens(torch.tensor([[tok]], device=mel.device))], dim=1)
-        return generated
-
-
-# Mel filterbank extracted from torchaudio (exact match, 0.0 diff)
-MEL_FB = torch.load(hf_hub_download("stukenov/sozkz-core-omniaudio-70m-kk-asr-v2", "mel_filterbank.pt", token=HF_TOKEN), map_location="cpu", weights_only=True)
+# Download and import original model code from HF repo
+model_code_path = hf_hub_download("stukenov/sozkz-core-omniaudio-70m-kk-asr-v1", "src/model_v2.py")
+import importlib.util
+spec = importlib.util.spec_from_file_location("model_v2", model_code_path)
+model_v2 = importlib.util.module_from_spec(spec)
+spec.loader.exec_module(model_v2)
+
+# Exact mel filterbank from torchaudio (pre-computed, diff=0.0)
+MEL_FB = torch.load(
+    hf_hub_download("stukenov/sozkz-core-omniaudio-70m-kk-asr-v2", "mel_filterbank.pt"),
+    map_location="cpu", weights_only=True,
+)
 MEL_WINDOW = torch.hann_window(400)
 
-def compute_mel(wav_np, sr=16000, n_fft=400, hop=160):
-    """Compute log-mel spectrogram matching torchaudio exactly."""
+def compute_mel(wav_np):
     wav = torch.from_numpy(wav_np).float()
-    stft = torch.stft(wav, n_fft=n_fft, hop_length=hop, win_length=n_fft,
+    stft = torch.stft(wav, n_fft=400, hop_length=160, win_length=400,
                       window=MEL_WINDOW, center=True, pad_mode="reflect", return_complex=True)
-    power = stft.abs().pow(2)  # (n_freqs, T)
-    mel = torch.matmul(MEL_FB.T, power)  # (80, T)
-    return torch.log(torch.clamp(mel, min=1e-10)).unsqueeze(0)  # (1, 80, T)
+    power = stft.abs().pow(2)
+    mel = torch.matmul(MEL_FB.T, power)
+    return torch.log(torch.clamp(mel, min=1e-10)).unsqueeze(0)
 
 
-# -- Load --
-
+# Load models
 ASR_MODELS = {
     "v2 (CTC+CE)": "stukenov/sozkz-core-omniaudio-70m-kk-asr-v2",
     "v1 (pure CE)": "stukenov/sozkz-core-omniaudio-70m-kk-asr-v1",
 }
-TOK_REPO = "stukenov/sozkz-core-gpt2-50k-kk-base-v1"
+ENC_CFG = {"n_mels": 80, "d_model": 256, "n_heads": 4, "n_layers": 6, "n_conv": 2}
+DEC_CFG = {"d_model": 512, "n_heads": 8, "n_layers": 8}
 
-print("Loading tokenizer...")
+TOK_REPO = "stukenov/sozkz-core-gpt2-50k-kk-base-v1"
 tok_file = hf_hub_download(TOK_REPO, "tokenizer.json")
 tokenizer = PreTrainedTokenizerFast(tokenizer_file=tok_file)
 tokenizer.eos_token = "<|endoftext|>"
 tokenizer.eos_token_id = 0
 
-def load_asr(repo):
-    print(f"Loading ASR from {repo}...")
-    model = OmniAudioScratchModel(
-        encoder_config={"n_mels": 80, "d_model": 256, "n_heads": 4, "n_layers": 6, "n_conv": 2},
-        decoder_config={"d_model": 512, "n_heads": 8, "n_layers": 8},
-        vocab_size=50257,
+loaded_asr = {}
+for name, repo in ASR_MODELS.items():
+    print(f"Loading {name} from {repo}...")
+    mdl = model_v2.OmniAudioScratchModel(
+        encoder_config=ENC_CFG, decoder_config=DEC_CFG, vocab_size=50257,
     )
     w = hf_hub_download(repo, "model.pt")
     sd = torch.load(w, map_location="cpu", weights_only=True)
-    missing, unexpected = model.load_state_dict(sd, strict=False)
-    model.lm_head.weight = model.embed_tokens.weight
-    model.requires_grad_(False)
-    params = sum(p.numel() for p in model.parameters()) / 1e6
-    print(f"  {params:.0f}M params, missing: {len(missing)}, unexpected: {len(unexpected)}")
-    return model
-
-loaded_asr = {}
-for name, repo in ASR_MODELS.items():
-    loaded_asr[name] = load_asr(repo)
-print("All ASR models loaded.")
+    info = mdl.load_state_dict(sd, strict=False)
+    print(f"  missing: {len(info.missing_keys)}, unexpected: {len(info.unexpected_keys)}")
+    for k in info.missing_keys:
+        if "rope" not in k and "inv_freq" not in k:
+            print(f"  MISSING: {k}")
+    mdl.requires_grad_(False)
+    loaded_asr[name] = mdl
+print("Ready.")
 
 
 @spaces.GPU
 def transcribe(audio, model_name):
-    import soundfile as sf
-
     if audio is None:
-        return "No audio provided"
+        return "No audio"
     t0 = time.perf_counter()
 
-    # Load audio as numpy float32 at 16kHz
+    # Load and resample to 16kHz mono
     if isinstance(audio, str):
         wav, sr = sf.read(audio)
         wav = np.array(wav, dtype=np.float32)
@@ -261,7 +94,7 @@ def transcribe(audio, model_name):
         if sr != 16000:
             wav = librosa.resample(wav, orig_sr=sr, target_sr=16000)
     else:
-        return "Unsupported audio format"
+        return "Unsupported format"
 
     wav = wav[:int(10.0 * 16000)]
     mel = compute_mel(wav)
@@ -314,7 +147,8 @@ with gr.Blocks(css=CSS, theme=theme, title="SozKZ ASR") as demo:
     gr.Markdown("Max 10 seconds. WAV/MP3/FLAC, 16kHz mono recommended.")
 
     gr.HTML("""<div style="text-align:center;padding:20px;font-size:12px;color:#aaa">
-        <a href="https://huggingface.co/stukenov/sozkz-core-omniaudio-70m-kk-asr-v1" style="color:#888">Model</a> |
+        <a href="https://huggingface.co/stukenov/sozkz-core-omniaudio-70m-kk-asr-v2" style="color:#888">v2 Model</a> |
+        <a href="https://huggingface.co/stukenov/sozkz-core-omniaudio-70m-kk-asr-v1" style="color:#888">v1 Model</a> |
         <a href="https://huggingface.co/spaces/stukenov/sozkz-kazakh-llm-demo" style="color:#888">LLM Demo</a> |
         <a href="https://huggingface.co/stukenov" style="color:#888">stukenov</a>
     </div>""")