Create run_eval_ml.py

run_eval_ml.py

@@ -0,0 +1,260 @@
+import argparse
+from typing import Optional
+import datasets
+from datasets import Audio
+import evaluate
+import soundfile as sf
+import tempfile
+import time
+import os
+import requests
+import itertools
+from tqdm import tqdm
+from dotenv import load_dotenv
+from normalizer import data_utils
+from normalizer.eval_utils import normalize_compound_pairs
+import concurrent.futures
+from providers import get_provider, PermanentError
+
+load_dotenv()
+
+
+def fetch_audio_urls(dataset_path, config_name, split, batch_size=100, max_retries=20):
+    API_URL = "https://datasets-server.huggingface.co/rows"
+
+    size_url = f"https://datasets-server.huggingface.co/size?dataset={dataset_path}&config={config_name}&split={split}"
+    size_response = requests.get(size_url).json()
+    total_rows = size_response["size"]["config"]["num_rows"]
+    for offset in tqdm(range(0, total_rows, batch_size), desc="Fetching audio URLs"):
+        params = {
+            "dataset": dataset_path,
+            "config": config_name,
+            "split": split,
+            "offset": offset,
+            "length": min(batch_size, total_rows - offset),
+        }
+
+        retries = 0
+        while retries <= max_retries:
+            try:
+                headers = {}
+                if os.environ.get("HF_TOKEN") is not None:
+                    headers["Authorization"] = f"Bearer {os.environ['HF_TOKEN']}"
+                else:
+                    print("HF_TOKEN not set, might experience rate-limiting.")
+                response = requests.get(API_URL, params=params)
+                response.raise_for_status()
+                data = response.json()
+                yield from data["rows"]
+                break
+            except (requests.exceptions.RequestException, ValueError) as e:
+                retries += 1
+                print(
+                    f"Error fetching data: {e}, retrying ({retries}/{max_retries})..."
+                )
+                time.sleep(10)
+                if retries >= max_retries:
+                    raise Exception("Max retries exceeded while fetching data.")
+
+
+def transcribe_with_retry(
+    model_name: str,
+    audio_file_path: Optional[str],
+    sample: dict,
+    max_retries=10,
+    use_url=False,
+    language="en",
+):
+    provider, variant = get_provider(model_name)
+    retries = 0
+    while retries <= max_retries:
+        try:
+            return provider.transcribe(variant, audio_file_path, sample, use_url=use_url, language=language)
+        except PermanentError:
+            raise
+        except Exception as e:
+            retries += 1
+            if retries > max_retries:
+                raise e
+
+            if not use_url:
+                sf.write(
+                    audio_file_path,
+                    sample["audio"]["array"],
+                    sample["audio"]["sampling_rate"],
+                    format="WAV",
+                )
+            delay = 1
+            print(
+                f"API Error: {str(e)}. Retrying in {delay}s... (Attempt {retries}/{max_retries})"
+            )
+            time.sleep(delay)
+
+
+def transcribe_dataset(
+    dataset_path,
+    config_name,
+    split,
+    model_name,
+    language,
+    use_url=False,
+    max_samples=None,
+    max_workers=4,
+):
+    if use_url:
+        audio_rows = fetch_audio_urls(dataset_path, config_name, split)
+        if max_samples:
+            audio_rows = itertools.islice(audio_rows, max_samples)
+        ds = audio_rows
+    else:
+        ds = datasets.load_dataset(dataset_path, config_name, split=split, streaming=False)
+        ds = ds.cast_column("audio", Audio(sampling_rate=16000))
+        if max_samples:
+            ds = ds.select(range(min(max_samples, len(ds))))
+
+    results = {
+        "references": [],
+        "predictions": [],
+        "audio_length_s": [],
+        "transcription_time_s": [],
+    }
+
+    print(f"Transcribing with model: {model_name}, language: {language}, config: {config_name}")
+
+    def process_sample(sample):
+        if use_url:
+            reference = sample["row"]["text"].strip()
+            audio_duration = sample["row"]["audio_length_s"]
+            start = time.time()
+            try:
+                transcription = transcribe_with_retry(
+                    model_name, None, sample, use_url=True, language=language
+                )
+            except Exception as e:
+                print(f"Failed to transcribe after retries: {e}")
+                return None
+
+        else:
+            reference = sample.get("text", "").strip()
+            with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmpfile:
+                sf.write(
+                    tmpfile.name,
+                    sample["audio"]["array"],
+                    sample["audio"]["sampling_rate"],
+                    format="WAV",
+                )
+                tmp_path = tmpfile.name
+                audio_duration = (
+                    len(sample["audio"]["array"]) / sample["audio"]["sampling_rate"]
+                )
+
+            start = time.time()
+            try:
+                transcription = transcribe_with_retry(
+                    model_name, tmp_path, sample, use_url=False, language=language
+                )
+            except Exception as e:
+                print(f"Failed to transcribe after retries: {e}")
+                os.unlink(tmp_path)
+                return None
+            finally:
+                if os.path.exists(tmp_path):
+                    os.unlink(tmp_path)
+                else:
+                    print(f"File {tmp_path} does not exist")
+
+        transcription_time = time.time() - start
+        return reference, transcription, audio_duration, transcription_time
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_sample = {
+            executor.submit(process_sample, sample): sample for sample in ds
+        }
+        for future in tqdm(
+            concurrent.futures.as_completed(future_to_sample),
+            total=len(future_to_sample),
+            desc="Transcribing",
+        ):
+            result = future.result()
+            if result:
+                reference, transcription, audio_duration, transcription_time = result
+                results["predictions"].append(transcription)
+                results["references"].append(reference)
+                results["audio_length_s"].append(audio_duration)
+                results["transcription_time_s"].append(transcription_time)
+
+    # Filter empty references (consistent with English pipeline's prepare_data)
+    filtered = [
+        (ref, pred, dur, time_s)
+        for ref, pred, dur, time_s in zip(
+            results["references"], results["predictions"],
+            results["audio_length_s"], results["transcription_time_s"]
+        )
+        if data_utils.is_target_text_in_range(ref)
+    ]
+    if filtered:
+        results["references"], results["predictions"], results["audio_length_s"], results["transcription_time_s"] = zip(*filtered)
+        results = {k: list(v) for k, v in results.items()}
+
+    manifest_path = data_utils.write_manifest(
+        results["references"],
+        results["predictions"],
+        model_name.replace("/", "-"),
+        dataset_path,
+        config_name,
+        split,
+        audio_length=results["audio_length_s"],
+        transcription_time=results["transcription_time_s"],
+    )
+
+    print("Results saved at path:", manifest_path)
+
+    norm_refs = [data_utils.ml_normalizer(r, lang=language) for r in results["references"]]
+    norm_preds = [data_utils.ml_normalizer(t, lang=language) for t in results["predictions"]]
+    wer_metric = evaluate.load("wer")
+    wer_refs, wer_preds = normalize_compound_pairs(norm_refs, norm_preds)
+    wer = wer_metric.compute(references=wer_refs, predictions=wer_preds)
+    wer_percent = round(100 * wer, 2)
+    rtfx = round(
+        sum(results["audio_length_s"]) / sum(results["transcription_time_s"]), 2
+    )
+
+    print("WER:", wer_percent, "%")
+    print("RTFx:", rtfx)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Multilingual API Transcription Script with Concurrency"
+    )
+    parser.add_argument("--dataset_path", required=True)
+    parser.add_argument("--config_name", required=True, help="Dataset config name, e.g. 'fleurs_de'")
+    parser.add_argument("--language", required=True, help="Language code, e.g. 'de'")
+    parser.add_argument("--split", default="test")
+    parser.add_argument(
+        "--model_name",
+        required=True,
+        help="Prefix model name with provider prefix (e.g., 'assembly/', 'openai/', 'elevenlabs/', 'revai/', 'speechmatics/' or 'aquavoice/')",
+    )
+    parser.add_argument("--max_samples", type=int, default=None)
+    parser.add_argument(
+        "--max_workers", type=int, default=300, help="Number of concurrent threads"
+    )
+    parser.add_argument(
+        "--use_url",
+        action="store_true",
+        help="Use URL-based audio fetching instead of datasets",
+    )
+
+    args = parser.parse_args()
+
+    transcribe_dataset(
+        dataset_path=args.dataset_path,
+        config_name=args.config_name,
+        split=args.split,
+        model_name=args.model_name,
+        language=args.language,
+        use_url=args.use_url,
+        max_samples=args.max_samples,
+        max_workers=args.max_workers,
+    )