admin-valentin/WhisperLiveKit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

translation now separates validated from output buffer tokens

Browse source
This commit is contained in:
Quentin Fuxa 2025-10-26 18:51:09 +01:00
parent 9434390ad3
commit 4e455b8aab
1 changed files with 26 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

50
whisperlivekit/translation/translation.py
View file

@ -35,17 +35,19 @@ class TranslationModel():
def load_model(src_langs, nllb_backend='ctranslate2', nllb_size='600M'): def load_model(src_langs, nllb_backend='ctranslate2', nllb_size='600M'):
device = "cuda" if torch.cuda.is_available() else "cpu" device = "cuda" if torch.cuda.is_available() else "cpu"
MODEL = f'nllb-200-distilled-{nllb_size}-ctranslate2'
if nllb_backend=='ctranslate2': if nllb_backend=='ctranslate2':
model = f'nllb-200-distilled-{nllb_size}-ctranslate2'
MODEL_GUY = 'entai2965' MODEL_GUY = 'entai2965'
huggingface_hub.snapshot_download(MODEL_GUY + '/' + MODEL,local_dir=MODEL) huggingface_hub.snapshot_download(MODEL_GUY + '/' + model,local_dir=model)
translator = ctranslate2.Translator(MODEL,device=device) translator = ctranslate2.Translator(model,device=device)
elif nllb_backend=='transformers': elif nllb_backend=='transformers':
translator = transformers.AutoModelForSeq2SeqLM.from_pretrained(f"facebook/nllb-200-distilled-{nllb_size}") model = f"facebook/nllb-200-distilled-{nllb_size}"
translator = transformers.AutoModelForSeq2SeqLM.from_pretrained(model)
tokenizer = dict() tokenizer = dict()
for src_lang in src_langs: for src_lang in src_langs:
if src_lang != 'auto': if src_lang != 'auto':
tokenizer[src_lang] = transformers.AutoTokenizer.from_pretrained(MODEL, src_lang=src_lang, clean_up_tokenization_spaces=True) tokenizer[src_lang] = transformers.AutoTokenizer.from_pretrained(model, src_lang=src_lang, clean_up_tokenization_spaces=True)
translation_model = TranslationModel( translation_model = TranslationModel(
translator=translator, translator=translator,
@ -61,7 +63,7 @@ def load_model(src_langs, nllb_backend='ctranslate2', nllb_size='600M'):
class OnlineTranslation: class OnlineTranslation:
def __init__(self, translation_model: TranslationModel, input_languages: list, output_languages: list): def __init__(self, translation_model: TranslationModel, input_languages: list, output_languages: list):
self.buffer = [] self.input_buffer = []
self.len_processed_buffer = 0 self.len_processed_buffer = 0
self.translation_remaining = Translation() self.translation_remaining = Translation()
self.validated = [] self.validated = []
@ -72,13 +74,13 @@ class OnlineTranslation:
def compute_common_prefix(self, results): def compute_common_prefix(self, results):
#we dont want want to prune the result for the moment. #we dont want want to prune the result for the moment.
if not self.buffer: if not self.input_buffer:
self.buffer = results self.input_buffer = results
else: else:
for i in range(min(len(self.buffer), len(results))): for i in range(min(len(self.input_buffer), len(results))):
if self.buffer[i] != results[i]: if self.input_buffer[i] != results[i]:
self.commited.extend(self.buffer[:i]) self.commited.extend(self.input_buffer[:i])
self.buffer = results[i:] self.input_buffer = results[i:]
def translate(self, input, input_lang, output_lang): def translate(self, input, input_lang, output_lang):
if not input: if not input:
@ -122,28 +124,28 @@ class OnlineTranslation:
def insert_tokens(self, tokens): def insert_tokens(self, tokens):
self.buffer.extend(tokens) self.input_buffer.extend(tokens)
pass pass
def process(self): def process(self):
i = 0 i = 0
if len(self.buffer) < self.len_processed_buffer + 3: #nothing new to process if len(self.input_buffer) < self.len_processed_buffer + 3: #nothing new to process
return self.validated + [self.translation_remaining] return self.validated + [self.translation_remaining]
while i < len(self.buffer): while i < len(self.input_buffer):
if self.buffer[i].is_punctuation(): if self.input_buffer[i].is_punctuation():
translation_sentence = self.translate_tokens(self.buffer[:i+1]) translation_sentence = self.translate_tokens(self.input_buffer[:i+1])
self.validated.append(translation_sentence) self.validated.append(translation_sentence)
self.buffer = self.buffer[i+1:] self.input_buffer = self.input_buffer[i+1:]
i = 0 i = 0
else: else:
i+=1 i+=1
self.translation_remaining = self.translate_tokens(self.buffer) self.translation_remaining = self.translate_tokens(self.input_buffer)
self.len_processed_buffer = len(self.buffer) self.len_processed_buffer = len(self.input_buffer)
return self.validated + [self.translation_remaining] return self.validated, [self.translation_remaining]
def insert_silence(self, silence_duration: float): def insert_silence(self, silence_duration: float):
if silence_duration >= MIN_SILENCE_DURATION_DEL_BUFFER: if silence_duration >= MIN_SILENCE_DURATION_DEL_BUFFER:
self.buffer = [] self.input_buffer = []
self.validated += [self.translation_remaining] self.validated += [self.translation_remaining]
if __name__ == '__main__': if __name__ == '__main__':
@ -157,13 +159,13 @@ if __name__ == '__main__':
test = test_string.split(' ') test = test_string.split(' ')
step = len(test) // 3 step = len(test) // 3
shared_model = load_model([input_lang], nllb_backend='ctranslate2') shared_model = load_model([input_lang], nllb_backend='transformers')
online_translation = OnlineTranslation(shared_model, input_languages=[input_lang], output_languages=[output_lang]) online_translation = OnlineTranslation(shared_model, input_languages=[input_lang], output_languages=[output_lang])
beg_inference = time.time() beg_inference = time.time()
for id in range(5): for id in range(5):
val = test[id*step : (id+1)*step] val = test[id*step : (id+1)*step]
val_str = ' '.join(val) val_str = ' '.join(val)
result = online_translation.translate(val_str) result = online_translation.translate(val_str, input_lang = input_lang, output_lang = output_lang)
print(result) print(result)
print('inference time:', time.time() - beg_inference) print('inference time:', time.time() - beg_inference)