Update app.py

app.py

@@ -1,5 +1,4 @@
-# app.py — Icelandic ASR (wav2vec2 primary, Whisper repair)
-# Dependencies: gradio, transformers, torch, librosa, soundfile
+# app.py — wav2vec2 multi-aug (stable + high quality)
 
 import os
 os.environ["OMP_NUM_THREADS"] = "1"
@@ -11,116 +10,65 @@ import torch
 import gc
 import re
 import librosa
-import numpy as np
 
 from transformers import (
     Wav2Vec2Processor,
-    Wav2Vec2ForCTC,
-    pipeline
+    Wav2Vec2ForCTC
 )
 
-# ============================================================
-# MODELS
-# ============================================================
-W2V_MODEL = "palli23/wav2vec2-icelandic-multi-aug-v2-5e-6"
-WHISPER_MODEL = "palli23/whisper-small-sam_spjall"
+#MODEL_ID = "palli23/wav2vec2-icelandic-multi-aug-v2-5e-6"
+MODEL_ID = "palli23/wav2vec2-icelandic-clean"
+# MODEL_ID = "palli23/wav2vec2-xlsr-300m-icelandic"
 
-# ============================================================
-# TEXT CLEANING (SAFE ONLY)
-# ============================================================
-def clean_ctc(text: str) -> str:
+# ——————————————————————————————
+# Strong Icelandic cleanup
+# ——————————————————————————————
+def clean_text(text: str) -> str:
     text = text.lower()
-    text = re.sub(r"(.)\1{3,}", r"\1\1", text)   # collapse char loops
-    text = re.sub(r"\s+", " ", text)
-    return text.strip()
-
-# ============================================================
-# STRUCTURAL VALIDITY CHECK (CRITICAL)
-# ============================================================
-def structurally_invalid(text: str) -> bool:
-    if not text or len(text) < 3:
-        return True
-
-    words = text.split()
-
-    # word-loop detection
-    for i in range(len(words) - 3):
-        if words[i:i+2] == words[i+2:i+4]:
-            return True
-
-    # collapsed word boundaries
-    if re.search(r"[a-záðéíóúýþæö]{18,}", text):
-        return True
-
-    # garbage consonant runs
-    if re.search(r"[bcdfghjklmnpqrstvwxz]{4,}", text):
-        return True
-
-    return False
 
-# ============================================================
-# OVERLAP MERGING (ANTI-DUPLICATION)
-# ============================================================
-def merge_chunks(prev: str, curr: str) -> str:
-    p = prev.split()
-    c = curr.split()
+    # collapse repeats (ctc artifacts)
+    text = re.sub(r"(.)\1{3,}", r"\1\1", text)
 
-    for k in range(6, 2, -1):
-        if len(p) >= k and len(c) >= k and p[-k:] == c[:k]:
-            return " ".join(p + c[k:])
-
-    return prev + " " + curr
-
-# ============================================================
-# SENTENCE FINALIZATION (LAST STEP ONLY)
-# ============================================================
-def finalize_text(text: str) -> str:
-    text = re.sub(r"\s+,", ",", text)
-    text = re.sub(r"\s+\.", ".", text)
-    text = re.sub(r"\s+\?", "?", text)
-    text = re.sub(r"\s+!", "!", text)
+    # spacing
+    text = re.sub(r"\s+", " ", text)
 
-    # light sentence segmentation
-    text = re.sub(r"([a-záðéíóúýþæö]) ([A-ZÁÉÍÓÚÝÞÆÖ])", r"\1.\n\2", text)
+    # punctuation spacing
+    text = text.replace(" ,", ",").replace(" .", ".")
+    text = text.replace(" ?", "?").replace(" !", "!")
 
     return text.strip()
 
-# ============================================================
-# AUDIO CHUNKING (STABLE)
-# ============================================================
+# ——————————————————————————————
+# Chunking helper (overlap improves WER)
+# ——————————————————————————————
 def chunk_audio(audio, sr, chunk_s=20, overlap_s=3):
-    step = int((chunk_s - overlap_s) * sr)
-    size = int(chunk_s * sr)
+    step = chunk_s - overlap_s
+    chunk_len = int(chunk_s * sr)
+    step_len = int(step * sr)
 
-    for start in range(0, len(audio), step):
-        chunk = audio[start:start + size]
-        if len(chunk) < sr:
+    for start in range(0, len(audio), step_len):
+        chunk = audio[start:start + chunk_len]
+        if len(chunk) < sr:  # too short
             break
         yield chunk
 
-# ============================================================
-# ZeroGPU WORKER
-# ============================================================
+# ——————————————————————————————
+# ZeroGPU worker
+# ——————————————————————————————
 @spaces.GPU(duration=180)
 def transcribe_3min(audio_path):
     if not audio_path:
         return "Hlaðið upp hljóðskrá"
 
-    # Load models
-    processor = Wav2Vec2Processor.from_pretrained(W2V_MODEL)
-    model = Wav2Vec2ForCTC.from_pretrained(W2V_MODEL).eval().to("cuda")
-
-    whisper = pipeline(
-        "automatic-speech-recognition",
-        model=WHISPER_MODEL,
-        device=0,
-        torch_dtype=torch.float16,
-    )
+    processor = Wav2Vec2Processor.from_pretrained(MODEL_ID)
+    model = Wav2Vec2ForCTC.from_pretrained(MODEL_ID)
+    model.eval().to("cuda")
 
+    # Load audio (float32 enforced)
     audio, sr = librosa.load(audio_path, sr=16000, mono=True)
     audio = audio.astype("float32")
 
-    outputs = []
+    texts = []
 
     for chunk in chunk_audio(audio, sr):
         inputs = processor(
@@ -131,51 +79,40 @@ def transcribe_3min(audio_path):
         )
 
         with torch.no_grad():
-            logits = model(inputs.input_values.to("cuda")).logits
+            logits = model(
+                inputs.input_values.to("cuda")
+            ).logits
 
         pred_ids = torch.argmax(logits, dim=-1)
-        w2v_text = clean_ctc(processor.batch_decode(pred_ids)[0])
+        text = processor.batch_decode(pred_ids)[0]
+        texts.append(text)
 
-        # -------- STRUCTURAL GATING --------
-        if structurally_invalid(w2v_text):
-            whisper_out = whisper(chunk, chunk_length_s=20)
-            text = whisper_out["text"].strip()
-        else:
-            text = w2v_text
+    final_text = clean_text(" ".join(texts))
 
-        # -------- MERGE WITH DEDUP --------
-        if outputs:
-            outputs[-1] = merge_chunks(outputs[-1], text)
-        else:
-            outputs.append(text)
-
-    final = finalize_text(" ".join(outputs))
-
-    # Cleanup
+    # Cleanup (critical)
     del model
     del processor
-    del whisper
     gc.collect()
     torch.cuda.empty_cache()
 
-    return final
+    return final_text
 
-# ============================================================
+# ——————————————————————————————
 # UI
-# ============================================================
+# ——————————————————————————————
 with gr.Blocks() as demo:
-    gr.Markdown("# 🇮🇸 Íslenskt ASR – Stable Hybrid (Fixed)")
-    gr.Markdown("**wav2vec2 multi-aug · structural gating · Whisper repair**")
+    gr.Markdown("# Íslenskt ASR – wav2vec2 (multi-aug)")
+    gr.Markdown("**stöðugt · chunked · post-processed**")
     gr.Markdown("**Hafa samband:** pallinr1@protonmail.com")
 
-    audio = gr.Audio(type="filepath", label="Hlaðið upp .wav / .mp3")
+    audio_in = gr.Audio(type="filepath", label="Hlaðið upp .mp3 / .wav")
     btn = gr.Button("Transcribe", variant="primary", size="lg")
-    out = gr.Textbox(lines=26, label="Útskrift")
+    output = gr.Textbox(lines=20, label="Útskrift")
 
-    btn.click(transcribe_3min, audio, out)
+    btn.click(fn=transcribe_3min, inputs=audio_in, outputs=output)
 
 demo.launch(
     share=True,
     server_name="0.0.0.0",
     server_port=7860,
-)
+)