File size: 3,396 Bytes
b9f09f7 |
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 |
from typing import Dict, Any, List
from datetime import timedelta
def format_time(seconds: float) -> str:
"""Format seconds as HH:MM:SS."""
return str(timedelta(seconds=int(seconds)))
def format_response(state: Dict[str, Any], with_diarization: bool = False) -> Dict[str, Any]:
"""
Format the shared state into a client-friendly response.
Args:
state: Current shared state dictionary
with_diarization: Whether to include diarization formatting
Returns:
Formatted response dictionary ready to send to client
"""
tokens = state["tokens"]
buffer_transcription = state["buffer_transcription"]
buffer_diarization = state["buffer_diarization"]
end_attributed_speaker = state["end_attributed_speaker"]
remaining_time_transcription = state["remaining_time_transcription"]
remaining_time_diarization = state["remaining_time_diarization"]
sep = state["sep"]
# Default response for empty state
if not tokens:
return {
"lines": [{
"speaker": 1,
"text": "",
"beg": format_time(0),
"end": format_time(0),
"diff": 0
}],
"buffer_transcription": buffer_transcription,
"buffer_diarization": buffer_diarization,
"remaining_time_transcription": remaining_time_transcription,
"remaining_time_diarization": remaining_time_diarization
}
# Process tokens to create response
previous_speaker = -1
lines = []
last_end_diarized = 0
undiarized_text = []
for token in tokens:
speaker = token.speaker
# Handle diarization logic
if with_diarization:
if (speaker == -1 or speaker == 0) and token.end >= end_attributed_speaker:
undiarized_text.append(token.text)
continue
elif (speaker == -1 or speaker == 0) and token.end < end_attributed_speaker:
speaker = previous_speaker
if speaker not in [-1, 0]:
last_end_diarized = max(token.end, last_end_diarized)
# Add new line or append to existing line
if speaker != previous_speaker or not lines:
lines.append({
"speaker": speaker,
"text": token.text,
"beg": format_time(token.start),
"end": format_time(token.end),
"diff": round(token.end - last_end_diarized, 2)
})
previous_speaker = speaker
elif token.text: # Only append if text isn't empty
lines[-1]["text"] += sep + token.text
lines[-1]["end"] = format_time(token.end)
lines[-1]["diff"] = round(token.end - last_end_diarized, 2)
# If we have undiarized text, include it in the buffer
if undiarized_text:
combined_buffer = sep.join(undiarized_text)
if buffer_transcription:
combined_buffer += sep + buffer_transcription
buffer_diarization = combined_buffer
return {
"lines": lines,
"buffer_transcription": buffer_transcription,
"buffer_diarization": buffer_diarization,
"remaining_time_transcription": remaining_time_transcription,
"remaining_time_diarization": remaining_time_diarization
} |