Improve the script

Author: dmontagu

pai-audio-evals/main.py

@@ -1,42 +1,44 @@
 import random
+import re
 from dataclasses import dataclass
+from functools import partial
 from pathlib import Path
 
 import logfire
+from nltk import edit_distance
 from pydantic import TypeAdapter
-from pydantic_ai import Agent, BinaryContent
+from pydantic_ai import Agent, BinaryContent, AudioUrl
+from pydantic_evals import Dataset, Case
+from pydantic_evals.evaluators import Evaluator, EvaluatorContext, EvaluatorOutput
 
-logfire.configure(service_name='pai-audio-evals')
-logfire.instrument_pydantic_ai()
-
-this_dir = Path(__file__).parent
-assets = this_dir / 'assets'
 
+@dataclass
+class EditSimilarity(Evaluator[object, str, object]):
+    def evaluate(self, ctx: EvaluatorContext[object, str, object]) -> EvaluatorOutput:
+        if ctx.expected_output is None:
+            return {}  # no metric
+        actual_tokens = re.sub(r'[^a-z0-9\s]', '', ctx.output.lower()).split()
+        expected_tokens = re.sub(r'[^a-z0-9\s]', '', ctx.expected_output.lower()).split()
+        distance = edit_distance(actual_tokens, expected_tokens)
+        normalized_distance = distance / max(len(actual_tokens), len(expected_tokens))
+        return 1 - normalized_distance
 
-def levenshtein_distance(s1: str, s2: str) -> int:
-    if len(s1) < len(s2):
-        return levenshtein_distance(s2, s1)
-    if len(s2) == 0:
-        return len(s1)
 
-    previous_row = range(len(s2) + 1)
-    for i, c1 in enumerate(s1):
-        current_row = [i + 1]
-        for j, c2 in enumerate(s2):
-            insertions = previous_row[j + 1] + 1
-            deletions = current_row[j] + 1
-            substitutions = previous_row[j] + (c1 != c2)
-            current_row.append(min(insertions, deletions, substitutions))
-        previous_row = current_row
+logfire.configure(service_name='pai-audio-evals', scrubbing=False, console=False)
+logfire.instrument_pydantic_ai()
 
-    return previous_row[-1]
+this_dir = Path(__file__).parent
+assets = this_dir / 'assets'
 
 
 @dataclass
 class AudioFile:
     file: str
     text: str
 
+    def audio_url(self) -> AudioUrl:
+        return AudioUrl(f'https://smokeshow.helpmanual.io/4l1l1s0s6q4741012x1w/{self.file}')
+
     def binary_content(self) -> BinaryContent:
         path = assets / self.file
         return BinaryContent(data=path.read_bytes(), media_type='audio/mpeg')
@@ -45,12 +47,18 @@ def binary_content(self) -> BinaryContent:
 files_schema = TypeAdapter(list[AudioFile])
 files = files_schema.validate_json((this_dir / 'assets.json').read_bytes())
 random.shuffle(files)
+
 audio_agent = Agent(instructions='return the transcription only, no prefix or quotes')
+dataset = Dataset(
+    cases=[Case(name=file.file, inputs=file.audio_url(), expected_output=file.text) for file in files],
+    evaluators=[EditSimilarity()],
+)
+
+
+async def task(audio_url: AudioUrl, model: str) -> str:
+    return (await audio_agent.run(['transcribe', audio_url], model=model)).output
+
 
-for audio_file in files[:3]:
-    with logfire.span('Transcribing audio {audio_file.text!r}', audio_file=audio_file):
-        model_distances: list[tuple[str, int]] = []
-        for model in 'gpt-4o-audio-preview', 'gpt-4o-mini-audio-preview', 'google-vertex:gemini-2.0-flash':
-            result = audio_agent.run_sync(['transcribe', audio_file.binary_content()], model=model)
-            model_distances.append((model, levenshtein_distance(audio_file.text, result.output)))
-        logfire.info(f'{model_distances}')
+with logfire.span('Compare models'):
+    for model in 'gpt-4o-audio-preview', 'gpt-4o-mini-audio-preview', 'google-vertex:gemini-2.0-flash':
+        dataset.evaluate_sync(partial(task, model=model), name=model, max_concurrency=10)

pyproject.toml

@@ -9,9 +9,13 @@ dependencies = [
     "devtools>=0.12.2",
     "fastapi>=0.115.14",
     "logfire[asyncpg,fastapi,httpx]>=3.21.1",
-    "pydantic-ai>=0.3.6",
+    "pydantic-ai>=0.4.5",
+    "nltk>=3.9.1",
 ]
 
+[tool.uv.sources]
+pydantic-ai = { git = "https://github.com/pydantic/pydantic-ai.git", rev = "0f46928bd07bc1a9f89c1d72c76cd2a86d52d489" }
+
 [dependency-groups]
 dev = ["ruff>=0.12.2", "asyncpg-stubs>=0.30.2"]