Evaluate Streaming transcription accuracy with WER

# pip install websocket-client jiwer whisper-normalizer
import websocket
import json
import os
import threading
import time
import wave
from urllib.parse import urlencode
import jiwer
from whisper_normalizer.basic import BasicTextNormalizer
from whisper_normalizer.english import EnglishTextNormalizer

# --- Configuration ---
ASSEMBLYAI_API_KEY = os.environ["ASSEMBLYAI_API_KEY"]
AUDIO_FILE = "audio.wav"  # Path to your audio file
SAMPLE_RATE = 48000  # Change to match the sample rate of your audio file

CONNECTION_PARAMS = {
    "speech_model": "u3-rt-pro",
    "sample_rate": SAMPLE_RATE,
}
API_ENDPOINT_BASE_URL = "wss://streaming.assemblyai.com/v3/ws"
API_ENDPOINT = f"{API_ENDPOINT_BASE_URL}?{urlencode(CONNECTION_PARAMS)}"

# Global variables
ws_app = None
audio_thread = None
stop_event = threading.Event()
assembly_streaming_transcript = ""

# --- WebSocket Event Handlers ---

def on_open(ws):
    """Called when the WebSocket connection is established."""
    print("WebSocket connection opened.")

    def stream_file():
        chunk_duration = 0.1

        with wave.open(AUDIO_FILE, 'rb') as wav_file:
            if wav_file.getnchannels() != 1:
                raise ValueError("Only mono audio is supported")

            file_sample_rate = wav_file.getframerate()
            if file_sample_rate != SAMPLE_RATE:
                print(f"Warning: File sample rate ({file_sample_rate}) doesn't match expected rate ({SAMPLE_RATE})")

            frames_per_chunk = int(file_sample_rate * chunk_duration)

            while not stop_event.is_set():
                frames = wav_file.readframes(frames_per_chunk)
                if not frames:
                    break
                ws.send(frames, websocket.ABNF.OPCODE_BINARY)

        print("File streaming complete. Waiting for final transcripts...")
        try:
            ws.send(json.dumps({"type": "Terminate"}))
        except Exception:
            pass

    global audio_thread
    audio_thread = threading.Thread(target=stream_file)
    audio_thread.daemon = True
    audio_thread.start()


def on_message(ws, message):
    global assembly_streaming_transcript
    try:
        data = json.loads(message)
        msg_type = data.get('type')

        if msg_type == "Begin":
            print(f"Session ID: {data.get('id')}")
        elif msg_type == "Turn":
            transcript = data.get('transcript', '')
            if data.get('end_of_turn'):
                assembly_streaming_transcript += transcript + " "
            print(transcript)
        elif msg_type == "Termination":
            print(f"Session terminated: {data.get('audio_duration_seconds', 0)} seconds of audio processed")
    except Exception as e:
        print(f"Error handling message: {e}")


def on_error(ws, error):
    """Called when a WebSocket error occurs."""
    print(f"\nWebSocket Error: {error}")
    stop_event.set()


def on_close(ws, close_status_code, close_msg):
    """Called when the WebSocket connection is closed."""
    print(f"\nWebSocket Disconnected: Status={close_status_code}")
    stop_event.set()
    if audio_thread and audio_thread.is_alive():
        audio_thread.join(timeout=1.0)


# --- Main Execution ---

ws_app = websocket.WebSocketApp(
    API_ENDPOINT,
    header={"Authorization": ASSEMBLYAI_API_KEY},
    on_open=on_open,
    on_message=on_message,
    on_error=on_error,
    on_close=on_close,
)

ws_thread = threading.Thread(target=ws_app.run_forever)
ws_thread.daemon = True
ws_thread.start()

try:
    while ws_thread.is_alive():
        time.sleep(0.1)
except KeyboardInterrupt:
    print("\nStopping...")
    stop_event.set()
    if ws_app:
        ws_app.close()
    ws_thread.join(timeout=2.0)

# --- Evaluate collected transcripts ---

reference_transcript = "AssemblyAI is a deep learning company that builds powerful APIs to help you transcribe and understand audio. The most common use case for the API is to automatically convert prerecorded audio and video files as well as real time audio streams into text transcriptions. Our APIs convert audio and video into text using powerful deep learning models that we research and develop end to end in house. Millions of podcasts, zoom recordings, phone calls or video files are being transcribed with Assembly AI every single day. But where Assembly AI really excels is with helping you understand your data. So let's say we transcribe Joe Biden's State of the Union using Assembly AI's API. With our Auto Chapters feature, you can generate time coded summaries of the key moments of your audio file. For example, with the State of the Union address we get chapter summaries like this. Auto Chapters automatically segments your audio or video files into chapters and provides a summary for each of these chapters. With Sentiment Analysis, we can classify what's being spoken in your audio files as either positive, negative or neutral. So for example, in the State of the Union address we see that this sentence was classified as positive, whereas this sentence was classified as negative. Content Safety Detection can flag sensitive content as it is spoken like hate speech, profanity, violence or weapons. For example, in Biden's State of the Union address, content safety detection flagged parts of his speech as being about weapons. This feature is especially useful for automatic content moderation and brand safety use cases. With Auto Highlights, you can automatically identify important words and phrases that are being spoken in your data owned by the State of the Union address. AssemblyAI's API detected these words and phrases as being important. Lastly, with entity detection you can identify entities that are spoken in your audio like organization names or person names. In Biden's speech, these were the entities that were detected. This is just a preview of the most popular features of AssemblyAI's API. If you want a full list of features, go check out our API documentation linked in the description below. And if you ever need some support, our team of developers is here to help. Everyday developers are using these features to build really exciting applications. From meeting summarizers to brand safety or contextual targeting platforms to full blown conversational intelligence tools. We can't wait to see what you build with AssemblyAI."

# Initialize normalizer
normalizer = EnglishTextNormalizer()
# For Spanish and other languages
# normalizer = BasicTextNormalizer()

def calculate_wer(reference, hypothesis, language='en'):
    # Normalize both texts
    normalized_reference = normalizer(reference)
    print("Reference: " + reference)
    print("Normalized Reference: " + normalized_reference + "\n")

    normalized_hypothesis = normalizer(hypothesis)
    print("Hypothesis: " + hypothesis)
    print("Normalized Hypothesis: " + normalized_hypothesis + "\n")

    # Calculate WER
    wer = jiwer.wer(normalized_reference, normalized_hypothesis)

    return wer * 100  # Return as percentage

wer_score = calculate_wer(reference_transcript, assembly_streaming_transcript.strip())
print(f"Final WER: {wer_score:.2f}%")

def on_open(ws):
    """Called when the WebSocket connection is established."""
    print("WebSocket connection opened.")

    def stream_file():
        chunk_duration = 0.1

        with wave.open(AUDIO_FILE, 'rb') as wav_file:
            if wav_file.getnchannels() != 1:
                raise ValueError("Only mono audio is supported")

            file_sample_rate = wav_file.getframerate()
            if file_sample_rate != SAMPLE_RATE:
                print(f"Warning: File sample rate ({file_sample_rate}) doesn't match expected rate ({SAMPLE_RATE})")

            frames_per_chunk = int(file_sample_rate * chunk_duration)

            while not stop_event.is_set():
                frames = wav_file.readframes(frames_per_chunk)
                if not frames:
                    break
                ws.send(frames, websocket.ABNF.OPCODE_BINARY)

        print("File streaming complete. Waiting for final transcripts...")
        try:
            ws.send(json.dumps({"type": "Terminate"}))
        except Exception:
            pass

    global audio_thread
    audio_thread = threading.Thread(target=stream_file)
    audio_thread.daemon = True
    audio_thread.start()


def on_message(ws, message):
    global assembly_streaming_transcript
    try:
        data = json.loads(message)
        msg_type = data.get('type')

        if msg_type == "Begin":
            print(f"Session ID: {data.get('id')}")
        elif msg_type == "Turn":
            transcript = data.get('transcript', '')
            if data.get('end_of_turn'):
                assembly_streaming_transcript += transcript + " "
            print(transcript)
        elif msg_type == "Termination":
            print(f"Session terminated: {data.get('audio_duration_seconds', 0)} seconds of audio processed")
    except Exception as e:
        print(f"Error handling message: {e}")


def on_error(ws, error):
    """Called when a WebSocket error occurs."""
    print(f"\nWebSocket Error: {error}")
    stop_event.set()


def on_close(ws, close_status_code, close_msg):
    """Called when the WebSocket connection is closed."""
    print(f"\nWebSocket Disconnected: Status={close_status_code}")
    stop_event.set()
    if audio_thread and audio_thread.is_alive():
        audio_thread.join(timeout=1.0)


# Connect and stream
ws_app = websocket.WebSocketApp(
    API_ENDPOINT,
    header={"Authorization": ASSEMBLYAI_API_KEY},
    on_open=on_open,
    on_message=on_message,
    on_error=on_error,
    on_close=on_close,
)

ws_thread = threading.Thread(target=ws_app.run_forever)
ws_thread.daemon = True
ws_thread.start()

try:
    while ws_thread.is_alive():
        time.sleep(0.1)
except KeyboardInterrupt:
    print("\nStopping...")
    stop_event.set()
    if ws_app:
        ws_app.close()
    ws_thread.join(timeout=2.0)