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:
@@ -2,6 +2,7 @@ import json
|
||||
import logging
|
||||
import threading
|
||||
import time
|
||||
import queue
|
||||
import numpy as np
|
||||
|
||||
|
||||
@@ -31,6 +32,7 @@ class ServeClientBase(object):
|
||||
no_speech_thresh=0.45,
|
||||
clip_audio=False,
|
||||
same_output_threshold=10,
|
||||
translation_queue=None,
|
||||
):
|
||||
self.client_uid = client_uid
|
||||
self.websocket = websocket
|
||||
@@ -50,6 +52,7 @@ class ServeClientBase(object):
|
||||
self.same_output_count = 0
|
||||
self.transcript = []
|
||||
self.end_time_for_same_output = None
|
||||
self.translation_queue = translation_queue
|
||||
|
||||
# threading
|
||||
self.lock = threading.Lock()
|
||||
@@ -307,7 +310,14 @@ class ServeClientBase(object):
|
||||
continue
|
||||
if self.get_segment_no_speech_prob(s) > self.no_speech_thresh:
|
||||
continue
|
||||
self.transcript.append(self.format_segment(start, end, text_, completed=True))
|
||||
completed_segment = self.format_segment(start, end, text_, completed=True)
|
||||
self.transcript.append(completed_segment)
|
||||
|
||||
if self.translation_queue:
|
||||
try:
|
||||
self.translation_queue.put(completed_segment.copy(), timeout=0.1)
|
||||
except queue.Full:
|
||||
logging.warning("Translation queue is full, skipping segment")
|
||||
offset = min(duration, self.get_segment_end(s))
|
||||
|
||||
# Process the last segment if its no_speech_prob is acceptable.
|
||||
@@ -340,12 +350,20 @@ class ServeClientBase(object):
|
||||
if not self.text or self.text[-1].strip().lower() != self.current_out.strip().lower():
|
||||
self.text.append(self.current_out)
|
||||
with self.lock:
|
||||
self.transcript.append(self.format_segment(
|
||||
completed_segment = self.format_segment(
|
||||
self.timestamp_offset,
|
||||
self.timestamp_offset + min(duration, self.end_time_for_same_output),
|
||||
self.current_out,
|
||||
completed=True
|
||||
))
|
||||
)
|
||||
self.transcript.append(completed_segment)
|
||||
|
||||
if self.translation_queue:
|
||||
try:
|
||||
self.translation_queue.put(completed_segment.copy(), timeout=0.1)
|
||||
except queue.Full:
|
||||
logging.warning("Translation queue is full, skipping segment")
|
||||
|
||||
self.current_out = ''
|
||||
offset = min(duration, self.end_time_for_same_output)
|
||||
self.same_output_count = 0
|
||||
|
||||
@@ -30,8 +30,9 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
send_last_n_segments=10,
|
||||
no_speech_thresh=0.45,
|
||||
clip_audio=False,
|
||||
same_output_threshold=10,
|
||||
cache_path="~/.cache/whisper-live/"
|
||||
same_output_threshold=7,
|
||||
cache_path="~/.cache/whisper-live/",
|
||||
translation_queue=None,
|
||||
):
|
||||
"""
|
||||
Initialize a ServeClient instance.
|
||||
@@ -61,6 +62,7 @@ class ServeClientFasterWhisper(ServeClientBase):
|
||||
no_speech_thresh,
|
||||
clip_audio,
|
||||
same_output_threshold,
|
||||
translation_queue
|
||||
)
|
||||
self.cache_path = cache_path
|
||||
self.model_sizes = [
|
||||
|
||||
@@ -0,0 +1,365 @@
|
||||
# Copyright (c) 2022 Idiap Research Institute, http://www.idiap.ch/
|
||||
# Written by Alireza Mohammadshahi <alireza.mohammadshahi@idiap.ch>
|
||||
# This is a modified version of https://github.com/huggingface/transformers/blob/main/src/transformers/models/m2m_100/tokenization_m2m_100.py
|
||||
# which owns by Fariseq Authors and The HuggingFace Inc. team.
|
||||
#
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
"""Tokenization classes for SMALL100."""
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from shutil import copyfile
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import sentencepiece
|
||||
|
||||
from transformers.tokenization_utils import BatchEncoding, PreTrainedTokenizer
|
||||
from transformers.utils import logging
|
||||
|
||||
|
||||
logger = logging.get_logger(__name__)
|
||||
|
||||
SPIECE_UNDERLINE = "▁"
|
||||
|
||||
VOCAB_FILES_NAMES = {
|
||||
"vocab_file": "vocab.json",
|
||||
"spm_file": "sentencepiece.bpe.model",
|
||||
"tokenizer_config_file": "tokenizer_config.json",
|
||||
}
|
||||
|
||||
PRETRAINED_VOCAB_FILES_MAP = {
|
||||
"vocab_file": {
|
||||
"alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/vocab.json",
|
||||
},
|
||||
"spm_file": {
|
||||
"alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/sentencepiece.bpe.model",
|
||||
},
|
||||
"tokenizer_config_file": {
|
||||
"alirezamsh/small100": "https://huggingface.co/alirezamsh/small100/resolve/main/tokenizer_config.json",
|
||||
},
|
||||
}
|
||||
|
||||
PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
|
||||
"alirezamsh/small100": 1024,
|
||||
}
|
||||
|
||||
# fmt: off
|
||||
FAIRSEQ_LANGUAGE_CODES = {
|
||||
"m2m100": ["af", "am", "ar", "ast", "az", "ba", "be", "bg", "bn", "br", "bs", "ca", "ceb", "cs", "cy", "da", "de", "el", "en", "es", "et", "fa", "ff", "fi", "fr", "fy", "ga", "gd", "gl", "gu", "ha", "he", "hi", "hr", "ht", "hu", "hy", "id", "ig", "ilo", "is", "it", "ja", "jv", "ka", "kk", "km", "kn", "ko", "lb", "lg", "ln", "lo", "lt", "lv", "mg", "mk", "ml", "mn", "mr", "ms", "my", "ne", "nl", "no", "ns", "oc", "or", "pa", "pl", "ps", "pt", "ro", "ru", "sd", "si", "sk", "sl", "so", "sq", "sr", "ss", "su", "sv", "sw", "ta", "th", "tl", "tn", "tr", "uk", "ur", "uz", "vi", "wo", "xh", "yi", "yo", "zh", "zu"]
|
||||
}
|
||||
# fmt: on
|
||||
|
||||
|
||||
class SMALL100Tokenizer(PreTrainedTokenizer):
|
||||
"""
|
||||
Construct an SMALL100 tokenizer. Based on [SentencePiece](https://github.com/google/sentencepiece).
|
||||
This tokenizer inherits from [`PreTrainedTokenizer`] which contains most of the main methods. Users should refer to
|
||||
this superclass for more information regarding those methods.
|
||||
Args:
|
||||
vocab_file (`str`):
|
||||
Path to the vocabulary file.
|
||||
spm_file (`str`):
|
||||
Path to [SentencePiece](https://github.com/google/sentencepiece) file (generally has a .spm extension) that
|
||||
contains the vocabulary.
|
||||
tgt_lang (`str`, *optional*):
|
||||
A string representing the target language.
|
||||
eos_token (`str`, *optional*, defaults to `"</s>"`):
|
||||
The end of sequence token.
|
||||
sep_token (`str`, *optional*, defaults to `"</s>"`):
|
||||
The separator token, which is used when building a sequence from multiple sequences, e.g. two sequences for
|
||||
sequence classification or for a text and a question for question answering. It is also used as the last
|
||||
token of a sequence built with special tokens.
|
||||
unk_token (`str`, *optional*, defaults to `"<unk>"`):
|
||||
The unknown token. A token that is not in the vocabulary cannot be converted to an ID and is set to be this
|
||||
token instead.
|
||||
pad_token (`str`, *optional*, defaults to `"<pad>"`):
|
||||
The token used for padding, for example when batching sequences of different lengths.
|
||||
language_codes (`str`, *optional*):
|
||||
What language codes to use. Should be `"m2m100"`.
|
||||
sp_model_kwargs (`dict`, *optional*):
|
||||
Will be passed to the `SentencePieceProcessor.__init__()` method. The [Python wrapper for
|
||||
SentencePiece](https://github.com/google/sentencepiece/tree/master/python) can be used, among other things,
|
||||
to set:
|
||||
- `enable_sampling`: Enable subword regularization.
|
||||
- `nbest_size`: Sampling parameters for unigram. Invalid for BPE-Dropout.
|
||||
- `nbest_size = {0,1}`: No sampling is performed.
|
||||
- `nbest_size > 1`: samples from the nbest_size results.
|
||||
- `nbest_size < 0`: assuming that nbest_size is infinite and samples from the all hypothesis (lattice)
|
||||
using forward-filtering-and-backward-sampling algorithm.
|
||||
- `alpha`: Smoothing parameter for unigram sampling, and dropout probability of merge operations for
|
||||
BPE-dropout.
|
||||
Examples:
|
||||
```python
|
||||
>>> from tokenization_small100 import SMALL100Tokenizer
|
||||
>>> tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100", tgt_lang="ro")
|
||||
>>> src_text = " UN Chief Says There Is No Military Solution in Syria"
|
||||
>>> tgt_text = "Şeful ONU declară că nu există o soluţie militară în Siria"
|
||||
>>> model_inputs = tokenizer(src_text, text_target=tgt_text, return_tensors="pt")
|
||||
>>> model(**model_inputs) # should work
|
||||
```"""
|
||||
|
||||
vocab_files_names = VOCAB_FILES_NAMES
|
||||
max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
|
||||
pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
|
||||
model_input_names = ["input_ids", "attention_mask"]
|
||||
|
||||
prefix_tokens: List[int] = []
|
||||
suffix_tokens: List[int] = []
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vocab_file,
|
||||
spm_file,
|
||||
tgt_lang=None,
|
||||
bos_token="<s>",
|
||||
eos_token="</s>",
|
||||
sep_token="</s>",
|
||||
pad_token="<pad>",
|
||||
unk_token="<unk>",
|
||||
language_codes="m2m100",
|
||||
sp_model_kwargs: Optional[Dict[str, Any]] = None,
|
||||
num_madeup_words=8,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
|
||||
|
||||
self.language_codes = language_codes
|
||||
fairseq_language_code = FAIRSEQ_LANGUAGE_CODES[language_codes]
|
||||
self.lang_code_to_token = {lang_code: f"__{lang_code}__" for lang_code in fairseq_language_code}
|
||||
|
||||
kwargs["additional_special_tokens"] = kwargs.get("additional_special_tokens", [])
|
||||
kwargs["additional_special_tokens"] += [
|
||||
self.get_lang_token(lang_code)
|
||||
for lang_code in fairseq_language_code
|
||||
if self.get_lang_token(lang_code) not in kwargs["additional_special_tokens"]
|
||||
]
|
||||
|
||||
self.vocab_file = vocab_file
|
||||
self.encoder = load_json(vocab_file)
|
||||
self.decoder = {v: k for k, v in self.encoder.items()}
|
||||
self.spm_file = spm_file
|
||||
self.sp_model = load_spm(spm_file, self.sp_model_kwargs)
|
||||
|
||||
self.encoder_size = len(self.encoder)
|
||||
|
||||
self.lang_token_to_id = {
|
||||
self.get_lang_token(lang_code): self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)
|
||||
}
|
||||
self.lang_code_to_id = {lang_code: self.encoder_size + i for i, lang_code in enumerate(fairseq_language_code)}
|
||||
self.id_to_lang_token = {v: k for k, v in self.lang_token_to_id.items()}
|
||||
|
||||
self._tgt_lang = tgt_lang if tgt_lang is not None else "en"
|
||||
self.cur_lang_id = self.get_lang_id(self._tgt_lang)
|
||||
self.num_madeup_words = num_madeup_words
|
||||
|
||||
super().__init__(
|
||||
tgt_lang=tgt_lang,
|
||||
bos_token=bos_token,
|
||||
eos_token=eos_token,
|
||||
sep_token=sep_token,
|
||||
unk_token=unk_token,
|
||||
pad_token=pad_token,
|
||||
language_codes=language_codes,
|
||||
sp_model_kwargs=self.sp_model_kwargs,
|
||||
num_madeup_words=num_madeup_words,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.set_lang_special_tokens(self._tgt_lang)
|
||||
|
||||
|
||||
@property
|
||||
def vocab_size(self) -> int:
|
||||
return len(self.encoder) + len(self.lang_token_to_id) + self.num_madeup_words
|
||||
|
||||
@property
|
||||
def tgt_lang(self) -> str:
|
||||
return self._tgt_lang
|
||||
|
||||
@tgt_lang.setter
|
||||
def tgt_lang(self, new_tgt_lang: str) -> None:
|
||||
self._tgt_lang = new_tgt_lang
|
||||
self.set_lang_special_tokens(self._tgt_lang)
|
||||
|
||||
def _tokenize(self, text: str) -> List[str]:
|
||||
return self.sp_model.encode(text, out_type=str)
|
||||
|
||||
def _convert_token_to_id(self, token):
|
||||
if token in self.lang_token_to_id:
|
||||
return self.lang_token_to_id[token]
|
||||
return self.encoder.get(token, self.encoder[self.unk_token])
|
||||
|
||||
def _convert_id_to_token(self, index: int) -> str:
|
||||
"""Converts an index (integer) in a token (str) using the decoder."""
|
||||
if index in self.id_to_lang_token:
|
||||
return self.id_to_lang_token[index]
|
||||
return self.decoder.get(index, self.unk_token)
|
||||
|
||||
def convert_tokens_to_string(self, tokens: List[str]) -> str:
|
||||
"""Converts a sequence of tokens (strings for sub-words) in a single string."""
|
||||
return self.sp_model.decode(tokens)
|
||||
|
||||
def get_special_tokens_mask(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
|
||||
) -> List[int]:
|
||||
"""
|
||||
Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
|
||||
special tokens using the tokenizer `prepare_for_model` method.
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
already_has_special_tokens (`bool`, *optional*, defaults to `False`):
|
||||
Whether or not the token list is already formatted with special tokens for the model.
|
||||
Returns:
|
||||
`List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
|
||||
"""
|
||||
|
||||
if already_has_special_tokens:
|
||||
return super().get_special_tokens_mask(
|
||||
token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
|
||||
)
|
||||
|
||||
prefix_ones = [1] * len(self.prefix_tokens)
|
||||
suffix_ones = [1] * len(self.suffix_tokens)
|
||||
if token_ids_1 is None:
|
||||
return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
|
||||
return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones
|
||||
|
||||
def build_inputs_with_special_tokens(
|
||||
self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
|
||||
) -> List[int]:
|
||||
"""
|
||||
Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and
|
||||
adding special tokens. An MBART sequence has the following format, where `X` represents the sequence:
|
||||
- `input_ids` (for encoder) `X [eos, src_lang_code]`
|
||||
- `decoder_input_ids`: (for decoder) `X [eos, tgt_lang_code]`
|
||||
BOS is never used. Pairs of sequences are not the expected use case, but they will be handled without a
|
||||
separator.
|
||||
Args:
|
||||
token_ids_0 (`List[int]`):
|
||||
List of IDs to which the special tokens will be added.
|
||||
token_ids_1 (`List[int]`, *optional*):
|
||||
Optional second list of IDs for sequence pairs.
|
||||
Returns:
|
||||
`List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens.
|
||||
"""
|
||||
if token_ids_1 is None:
|
||||
if self.prefix_tokens is None:
|
||||
return token_ids_0 + self.suffix_tokens
|
||||
else:
|
||||
return self.prefix_tokens + token_ids_0 + self.suffix_tokens
|
||||
# We don't expect to process pairs, but leave the pair logic for API consistency
|
||||
if self.prefix_tokens is None:
|
||||
return token_ids_0 + token_ids_1 + self.suffix_tokens
|
||||
else:
|
||||
return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
|
||||
|
||||
def get_vocab(self) -> Dict:
|
||||
vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
|
||||
vocab.update(self.added_tokens_encoder)
|
||||
return vocab
|
||||
|
||||
def __getstate__(self) -> Dict:
|
||||
state = self.__dict__.copy()
|
||||
state["sp_model"] = None
|
||||
return state
|
||||
|
||||
def __setstate__(self, d: Dict) -> None:
|
||||
self.__dict__ = d
|
||||
|
||||
# for backward compatibility
|
||||
if not hasattr(self, "sp_model_kwargs"):
|
||||
self.sp_model_kwargs = {}
|
||||
|
||||
self.sp_model = load_spm(self.spm_file, self.sp_model_kwargs)
|
||||
|
||||
def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
|
||||
save_dir = Path(save_directory)
|
||||
if not save_dir.is_dir():
|
||||
raise OSError(f"{save_directory} should be a directory")
|
||||
vocab_save_path = save_dir / (
|
||||
(filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]
|
||||
)
|
||||
spm_save_path = save_dir / (
|
||||
(filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["spm_file"]
|
||||
)
|
||||
|
||||
save_json(self.encoder, vocab_save_path)
|
||||
|
||||
if os.path.abspath(self.spm_file) != os.path.abspath(spm_save_path) and os.path.isfile(self.spm_file):
|
||||
copyfile(self.spm_file, spm_save_path)
|
||||
elif not os.path.isfile(self.spm_file):
|
||||
with open(spm_save_path, "wb") as fi:
|
||||
content_spiece_model = self.sp_model.serialized_model_proto()
|
||||
fi.write(content_spiece_model)
|
||||
|
||||
return (str(vocab_save_path), str(spm_save_path))
|
||||
|
||||
def prepare_seq2seq_batch(
|
||||
self,
|
||||
src_texts: List[str],
|
||||
tgt_texts: Optional[List[str]] = None,
|
||||
tgt_lang: str = "ro",
|
||||
**kwargs,
|
||||
) -> BatchEncoding:
|
||||
self.tgt_lang = tgt_lang
|
||||
self.set_lang_special_tokens(self.tgt_lang)
|
||||
return super().prepare_seq2seq_batch(src_texts, tgt_texts, **kwargs)
|
||||
|
||||
def _build_translation_inputs(self, raw_inputs, tgt_lang: Optional[str], **extra_kwargs):
|
||||
"""Used by translation pipeline, to prepare inputs for the generate function"""
|
||||
if tgt_lang is None:
|
||||
raise ValueError("Translation requires a `tgt_lang` for this model")
|
||||
self.tgt_lang = tgt_lang
|
||||
inputs = self(raw_inputs, add_special_tokens=True, **extra_kwargs)
|
||||
return inputs
|
||||
|
||||
def _switch_to_input_mode(self):
|
||||
self.set_lang_special_tokens(self.tgt_lang)
|
||||
|
||||
def _switch_to_target_mode(self):
|
||||
self.prefix_tokens = None
|
||||
self.suffix_tokens = [self.eos_token_id]
|
||||
|
||||
def set_lang_special_tokens(self, src_lang: str) -> None:
|
||||
"""Reset the special tokens to the tgt lang setting. No prefix and suffix=[eos, tgt_lang_code]."""
|
||||
lang_token = self.get_lang_token(src_lang)
|
||||
self.cur_lang_id = self.lang_token_to_id[lang_token]
|
||||
self.prefix_tokens = [self.cur_lang_id]
|
||||
self.suffix_tokens = [self.eos_token_id]
|
||||
|
||||
def get_lang_token(self, lang: str) -> str:
|
||||
return self.lang_code_to_token[lang]
|
||||
|
||||
def get_lang_id(self, lang: str) -> int:
|
||||
lang_token = self.get_lang_token(lang)
|
||||
return self.lang_token_to_id[lang_token]
|
||||
|
||||
|
||||
def load_spm(path: str, sp_model_kwargs: Dict[str, Any]) -> sentencepiece.SentencePieceProcessor:
|
||||
spm = sentencepiece.SentencePieceProcessor(**sp_model_kwargs)
|
||||
spm.Load(str(path))
|
||||
return spm
|
||||
|
||||
|
||||
def load_json(path: str) -> Union[Dict, List]:
|
||||
with open(path, "r") as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def save_json(data, path: str) -> None:
|
||||
with open(path, "w") as f:
|
||||
json.dump(data, f, indent=2)
|
||||
@@ -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