Spaces:
Runtime error
Runtime error
File size: 4,582 Bytes
1207440 | 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 | import asyncio
import websockets
import numpy as np
from .whisper_vad import RealtimeWhisperVAD
import json
import logging
import time
import aiohttp
logger = logging.getLogger(__name__)
clients = set()
# Simple transcript assembly utilities
class TranscriptBuffer:
def __init__(self, silence_threshold=1.0):
self.words = []
self.last_word_time = None
self.silence_threshold = silence_threshold
def add_word(self, word_obj):
self.words.append(word_obj.word)
self.last_word_time = time.time()
def should_flush(self):
if not self.words:
return False
return (time.time() - (self.last_word_time or 0)) > self.silence_threshold
def flush(self):
txt = " ".join(self.words).strip()
self.words = []
self.last_word_time = None
return txt
async def post_transcript_to_rag(transcript: str):
url = "http://localhost:5001/evaluate"
payload = {"transcript": transcript}
try:
async with aiohttp.ClientSession() as session:
async with session.post(url, json=payload, timeout=10) as resp:
text = await resp.text()
logger.info(f"Posted transcript to RAG; response status={resp.status}")
return resp.status, text
except Exception as e:
logger.error(f"Failed to post transcript to RAG: {e}")
return None, str(e)
# Callback function for RealtimeWhisperVAD to call per word
transcript_buffer = TranscriptBuffer(silence_threshold=1.0)
def word_callback(word):
# This callback runs in a thread; schedule async broadcast and flush checks
try:
# Add to transcript buffer
transcript_buffer.add_word(word)
# Broadcast word to connected clients asynchronously
msg = json.dumps({"word": word.word, "start": word.start, "end": word.end})
asyncio.run_coroutine_threadsafe(broadcast_message(msg), asyncio.get_event_loop())
except Exception as e:
logger.error(f"Error in word_callback: {e}")
async def broadcast_message(message: str):
dead = []
for client in clients:
try:
await client.send(message)
except Exception:
dead.append(client)
for d in dead:
clients.discard(d)
# Configuration for RealtimeWhisperVAD
config = {
"whisper_model": "tiny",
"sample_rate": 16000,
"chunk_duration": 0.9,
"vad_threshold": 0.5,
"min_speech_duration": 0.5,
"max_speech_duration": 5,
"silence_duration": 0.5,
"transcription_callback": word_callback,
}
# Initialize RealtimeWhisperVAD
transcriber = RealtimeWhisperVAD(**config)
async def handler(websocket, path):
logger.info(f"Client connected from {websocket.remote_address}")
clients.add(websocket)
try:
# If STT backend is unavailable, inform client and close
if not getattr(transcriber, 'use_faster_whisper', False):
await websocket.send(json.dumps({"error": "STT backend unavailable"}))
await websocket.close()
return
# Start the transcriber when a client connects
if not transcriber.running:
transcriber.start()
# Polling loop: receive audio chunks and periodically check buffer flush
async for message in websocket:
try:
audio_float32 = np.frombuffer(message, dtype=np.float32)
transcriber.add_audio_chunk(audio_float32)
except Exception as e:
logger.error(f"Error processing audio chunk: {e}")
# If buffer indicates silence, flush transcript and POST to RAG
if transcript_buffer.should_flush():
transcript = transcript_buffer.flush()
if transcript:
status, resp_text = await post_transcript_to_rag(transcript)
# Send evaluation result back to clients
payload = {"event": "evaluation_result", "status": status, "resp": resp_text}
await broadcast_message(json.dumps(payload))
except websockets.exceptions.ConnectionClosed:
logger.info(f"Client disconnected from {websocket.remote_address}")
finally:
clients.discard(websocket)
async def main():
# Use 0.0.0.0 to be reachable; keep localhost if you want local-only
async with websockets.serve(handler, "0.0.0.0", 5002):
logger.info("WebSocket server started on ws://0.0.0.0:5002")
await asyncio.Future() # run forever
if __name__ == "__main__":
asyncio.run(main()) |