File size: 5,751 Bytes
ba4a241
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
030e33b
 
 
 
 
ba4a241
 
030e33b
ba4a241
 
 
 
 
 
030e33b
ba4a241
 
 
 
 
 
030e33b
 
 
 
 
 
 
 
 
 
 
ba4a241
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
030e33b
ba4a241
030e33b
 
ba4a241
 
 
 
 
 
030e33b
 
ba4a241
 
 
 
 
 
 
 
 
 
 
 
 
 
 
030e33b
ba4a241
 
 
 
 
 
 
 
030e33b
 
 
 
 
 
 
 
 
 
 
 
ba4a241
030e33b
 
ba4a241
 
 
 
 
030e33b
ba4a241
 
 
 
 
030e33b
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
from __future__ import annotations

from pathlib import Path
from typing import Dict, Iterable, List, Optional, Tuple

import soundfile as sf
from fastapi import HTTPException

from src.asr import transcribe_file
from src.diarization import (
    get_diarization_stats,
    init_speaker_embedding_extractor,
    merge_consecutive_utterances,
    merge_transcription_with_diarization,
    perform_speaker_diarization_on_utterances,
)
from src.utils import sensevoice_models

from ..core.config import get_settings
from ..models.transcription import DiarizationOptions, TranscriptionRequest

settings = get_settings()


def _serialize_utterance(utt: Tuple[float, float, str], speaker: Optional[int] = None) -> Dict[str, object]:
    start, end, text = utt
    payload: Dict[str, object] = {
        "start": round(float(start), 3),
        "end": round(float(end), 3),
        "text": text,
    }
    if speaker is not None:
        payload["speaker"] = int(speaker)
    return payload


def _prepare_model_name(options: TranscriptionRequest) -> str:
    if options.backend == "sensevoice":
        # sensevoice_models stores map from friendly name to repo id
        return sensevoice_models.get(options.model_name, options.model_name)
    return options.model_name


def iter_transcription_events(
    audio_path: Path,
    audio_url: str,
    options: TranscriptionRequest,
) -> Iterable[Dict[str, object]]:
    model_name = _prepare_model_name(options)

    try:
        generator = transcribe_file(
            audio_path=str(audio_path),
            vad_threshold=options.vad_threshold,
            model_name=model_name,
            backend=options.backend,
            language=options.language,
            textnorm=options.textnorm,
        )

        yield {
            "type": "ready",
            "audioUrl": audio_url,
            "backend": options.backend,
            "model": model_name,
        }

        yield {
            "type": "status",
            "message": "Transcribing audio...",
        }

        final_utterances: List[Tuple[float, float, str]] = []

        for current_utterance, all_utterances, progress in generator:
            if current_utterance:
                start, end, text = current_utterance
                yield {
                    "type": "utterance",
                    "utterance": _serialize_utterance((start, end, text)),
                    "index": len(all_utterances) - 1,
                    "progress": round(progress, 1),
                }
            final_utterances = list(all_utterances)

        # Final event with transcript and optional diarization
        diarization_payload = None
        if options.diarization.enable:
            yield {
                "type": "status",
                "message": "Performing speaker diarization...",
            }
            diarization_gen = _run_diarization(audio_path, final_utterances, options.diarization)
            for event in diarization_gen:
                if event["type"] == "progress":
                    yield event
                elif event["type"] == "result":
                    diarization_payload = event["payload"]
                    break

        transcript_text = "\n".join([utt[2] for utt in final_utterances])

        yield {
            "type": "complete",
            "utterances": [_serialize_utterance(utt) for utt in final_utterances],
            "transcript": transcript_text,
            "diarization": diarization_payload,
        }

    except Exception as exc:  # pragma: no cover
        raise HTTPException(status_code=500, detail=f"Transcription failed: {exc}")


def _run_diarization(
    audio_path: Path,
    utterances: List[Tuple[float, float, str]],
    options: DiarizationOptions,
):
    if not utterances:
        yield {"type": "result", "payload": None}
        return

    extractor_result = init_speaker_embedding_extractor(
        cluster_threshold=options.cluster_threshold,
        num_speakers=options.num_speakers,
    )
    if not extractor_result:
        yield {"type": "result", "payload": None}
        return

    embedding_extractor, config_dict = extractor_result

    audio, sample_rate = sf.read(str(audio_path), dtype="float32")
    if audio.ndim > 1:
        audio = audio.mean(axis=1)

    if sample_rate != 16000:
        # Lazy import to avoid mandatory dependency during module import
        from scipy.signal import resample

        target_num_samples = int(len(audio) * 16000 / sample_rate)
        audio = resample(audio, target_num_samples)
        sample_rate = 16000

    diarization_gen = perform_speaker_diarization_on_utterances(
        audio=audio,
        sample_rate=sample_rate,
        utterances=utterances,
        embedding_extractor=embedding_extractor,
        config_dict=config_dict,
        progress_callback=None,
    )

    diarization_segments = None
    try:
        while True:
            item = next(diarization_gen)
            if isinstance(item, float):
                yield {"type": "progress", "stage": "diarization", "progress": round(item * 100, 1)}
            else:
                diarization_segments = item
                break
    except StopIteration as e:
        diarization_segments = e.value

    if not diarization_segments:
        yield {"type": "result", "payload": None}
        return

    merged = merge_transcription_with_diarization(utterances, diarization_segments)
    merged = merge_consecutive_utterances(merged, max_gap=1.0)
    stats = get_diarization_stats(merged)

    yield {"type": "result", "payload": {
        "utterances": [
            _serialize_utterance((start, end, text), speaker)
            for start, end, text, speaker in merged
        ],
        "stats": stats,
    }}