Add translation backend
Translate from any language to any language with alirezamsh/small100 running in a thread and reading from a queue shared with transcription thread. Signed-off-by: makaveli10 <vineet.suryan@collabora.com>
This commit is contained in:
makaveli10
@@ -0,0 +1,218 @@
|
||||
import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import queue
|
||||
from typing import Dict, Any, Optional
|
||||
import torch
|
||||
import threading
|
||||
from transformers import M2M100ForConditionalGeneration
|
||||
from whisper_live.backend.tokenization_small100 import SMALL100Tokenizer
|
||||
|
||||
from whisper_live.backend.base import ServeClientBase
|
||||
|
||||
|
||||
class ServeClientTranslation(ServeClientBase):
|
||||
"""
|
||||
Handles translation of completed transcription segments in a separate thread.
|
||||
Reads from a queue populated by the transcription backend and sends translated
|
||||
segments back to the client via WebSocket.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client_uid,
|
||||
websocket,
|
||||
translation_queue,
|
||||
target_language="fr",
|
||||
send_last_n_segments=10,
|
||||
model_name="alirezamsh/small100"
|
||||
):
|
||||
"""
|
||||
Initialize the translation client.
|
||||
|
||||
Args:
|
||||
client_uid (str): Unique identifier for the client
|
||||
websocket: WebSocket connection to the client
|
||||
translation_queue (queue.Queue): Queue containing completed segments to translate
|
||||
target_language (str): Target language code (default: "fr" for French)
|
||||
send_last_n_segments (int): Number of recent translated segments to send
|
||||
model_name (str): Translation model name to use
|
||||
"""
|
||||
super().__init__(client_uid, websocket, send_last_n_segments)
|
||||
self.translation_queue = translation_queue
|
||||
self.target_language = target_language
|
||||
self.model_name = model_name
|
||||
self.translated_segments = []
|
||||
self.translation_model = None
|
||||
self.tokenizer = None
|
||||
self.device = None
|
||||
self.model_loaded = False
|
||||
self.load_translation_model()
|
||||
|
||||
def load_translation_model(self):
|
||||
"""Load the translation model and tokenizer."""
|
||||
try:
|
||||
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
logging.info(f"Loading translation model on device: {self.device}")
|
||||
|
||||
self.translation_model = M2M100ForConditionalGeneration.from_pretrained(
|
||||
self.model_name
|
||||
).to(self.device)
|
||||
self.tokenizer = SMALL100Tokenizer.from_pretrained(self.model_name)
|
||||
self.tokenizer.tgt_lang = self.target_language
|
||||
|
||||
self.model_loaded = True
|
||||
logging.info(f"Translation model loaded successfully. Target language: {self.target_language}")
|
||||
except Exception as e:
|
||||
logging.error(f"Failed to load translation model: {e}")
|
||||
self.translation_model = None
|
||||
self.tokenizer = None
|
||||
self.model_loaded = False
|
||||
|
||||
def translate_text(self, text: str) -> str:
|
||||
"""
|
||||
Translate a single text segment.
|
||||
|
||||
Args:
|
||||
text (str): Text to translate
|
||||
|
||||
Returns:
|
||||
str: Translated text or original text if translation fails
|
||||
"""
|
||||
if not self.model_loaded or not text.strip():
|
||||
return text
|
||||
|
||||
try:
|
||||
# Encode input and move to device
|
||||
encoded_input = self.tokenizer(text, return_tensors="pt").to(self.device)
|
||||
|
||||
# Generate translation
|
||||
with torch.no_grad():
|
||||
generated_tokens = self.translation_model.generate(**encoded_input)
|
||||
|
||||
# Decode output
|
||||
output = self.tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
|
||||
return output[0] if output else text
|
||||
|
||||
except Exception as e:
|
||||
logging.error(f"Translation failed for text '{text}': {e}")
|
||||
return text
|
||||
|
||||
def process_translation_queue(self):
|
||||
"""
|
||||
Process segments from the translation queue.
|
||||
Continuously reads from the queue until None is received (exit signal).
|
||||
"""
|
||||
logging.info(f"Starting translation processing for client {self.client_uid}")
|
||||
|
||||
while not self.exit:
|
||||
try:
|
||||
# Get segment from queue with timeout
|
||||
segment = self.translation_queue.get(timeout=1.0)
|
||||
|
||||
# Check for exit signal
|
||||
if segment is None:
|
||||
logging.info(f"Received exit signal for translation client {self.client_uid}")
|
||||
break
|
||||
|
||||
# Only translate completed segments
|
||||
if not segment.get("completed", False):
|
||||
self.translation_queue.task_done()
|
||||
continue
|
||||
|
||||
# Translate the segment
|
||||
original_text = segment.get("text", "")
|
||||
translated_text = self.translate_text(original_text)
|
||||
|
||||
# Create translated segment
|
||||
translated_segment = {
|
||||
"start": segment["start"],
|
||||
"end": segment["end"],
|
||||
"text": translated_text,
|
||||
"completed": segment.get("completed", False),
|
||||
"target_language": self.target_language
|
||||
}
|
||||
|
||||
self.translated_segments.append(translated_segment)
|
||||
segments_to_send = self.prepare_translated_segments()
|
||||
self.send_translation_to_client(segments_to_send)
|
||||
|
||||
self.translation_queue.task_done()
|
||||
|
||||
except queue.Empty:
|
||||
continue
|
||||
except Exception as e:
|
||||
logging.error(f"Error processing translation queue: {e}")
|
||||
continue
|
||||
|
||||
logging.info(f"Translation processing ended for client {self.client_uid}")
|
||||
|
||||
def prepare_translated_segments(self):
|
||||
"""
|
||||
Prepare the last n translated segments to send to client.
|
||||
|
||||
Returns:
|
||||
list: List of recent translated segments
|
||||
"""
|
||||
if len(self.translated_segments) >= self.send_last_n_segments:
|
||||
return self.translated_segments[-self.send_last_n_segments:]
|
||||
return self.translated_segments[:]
|
||||
|
||||
def send_translation_to_client(self, translated_segments):
|
||||
"""
|
||||
Send translated segments to the client via WebSocket.
|
||||
|
||||
Args:
|
||||
translated_segments (list): List of translated segments to send
|
||||
"""
|
||||
try:
|
||||
self.websocket.send(
|
||||
json.dumps({
|
||||
"uid": self.client_uid,
|
||||
"translated_segments": translated_segments,
|
||||
})
|
||||
)
|
||||
except Exception as e:
|
||||
logging.error(f"[ERROR]: Sending translation data to client: {e}")
|
||||
|
||||
def speech_to_text(self):
|
||||
"""
|
||||
Override parent method to handle translation processing.
|
||||
This method will be called when the translation thread starts.
|
||||
"""
|
||||
self.process_translation_queue()
|
||||
|
||||
def set_target_language(self, language: str):
|
||||
"""
|
||||
Change the target language for translation.
|
||||
|
||||
Args:
|
||||
language (str): New target language code
|
||||
"""
|
||||
self.target_language = language
|
||||
if self.tokenizer:
|
||||
self.tokenizer.tgt_lang = language
|
||||
logging.info(f"Target language changed to: {language}")
|
||||
|
||||
def cleanup(self):
|
||||
"""Clean up translation resources."""
|
||||
logging.info(f"Cleaning up translation resources for client {self.client_uid}")
|
||||
self.exit = True
|
||||
|
||||
try:
|
||||
self.translation_queue.put(None, timeout=1.0)
|
||||
except:
|
||||
pass
|
||||
|
||||
self.translated_segments.clear()
|
||||
|
||||
if self.translation_model:
|
||||
del self.translation_model
|
||||
self.translation_model = None
|
||||
if self.tokenizer:
|
||||
del self.tokenizer
|
||||
self.tokenizer = None
|
||||
|
||||
if self.device and self.device.type == 'cuda':
|
||||
torch.cuda.empty_cache()
|
||||
Reference in New Issue
Block a user