whispertube_backend / translation.py
uzi007's picture
APIs updated & yt-dlp added
fe9dbf9
raw
history blame
4.09 kB
import os
import openai
import locale
locale.getpreferredencoding = lambda: "UTF-8"
import dl_translate as dlt
from deep_translator import GoogleTranslator
from languages import LANGUAGES
OPENAI_API_KEY = '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'
OPENAI_API_URL = 'https://api.openai.com/v1/chat/completions'
openai.api_key = OPENAI_API_KEY
class Translation:
def __init__(self, transcript_dict, source_lang, target_lang, output_path):
self.transcript_dict = transcript_dict
self.output_path = os.path.join(os.getcwd(), output_path)
# Languages
self.source_lang = source_lang # Whisper Detected Language
self.target_lang = target_lang
# Transcript
self.transcript = transcript_dict['text'].strip()
self.subtitles = self.__get_subtitles()
# Translation Model
nllb_model = 'facebook/nllb-200-distilled-600M'
# nllb_model = 'facebook/nllb-200-1.3B'
# nllb_model = 'facebook/nllb-200-3.3B'
# nllb_model = 'facebook/nllb-moe-54b'
self.nllb = dlt.TranslationModel(nllb_model)
def __get_subtitles(self):
'''
Returns the subtitles from transcript dictionary
'''
subtitles = []
for s in self.transcript_dict['segments']:
segment = {
'start': s['start'],
'end': s['end'],
'text': s['text'].strip()
}
subtitles.append(segment)
return subtitles
def __correct_punctuation_gpt(self):
'''
Corrects the Punctuation from GPT
'''
system_prompt = """
You are a helpful NLP assistant.
Your task is to identify language of the provided text,
correct any spelling discrepancies in the transcribed text
as well as add punctuation in the multilingual text if they are missing.
Only add necessary punctuation such as periods, commas, and capitalization,
and use only the context provided.
You response should be as follows:
Corrected Text:
Here goes the corrected text with punctuation.
"""
user_prompt = f"""
Here is the text:
{self.transcript}
"""
response = openai.ChatCompletion.create(
model="gpt-3.5-turbo",
messages=[
{"role": "system", "content": system_prompt},
{"role": "user", "content": user_prompt},
]
)
text = response.choices[0].message.content.replace('Corrected Text:\n', '')
return text
def get_translated_transcript(self):
'''
Translates the transcript into required language
'''
# Correcting Punctuation using GPT
transcript = self.__correct_punctuation_gpt()
# Splitting Text into Sentences
if self.source_lang in ['ar', 'ur']:
splitter = '۔'
else:
splitter = '.'
sentences = transcript.split(splitter)
# Getting Translation using NLLB
translated_transcript = ''
for sentence in sentences:
translated_sentence = self.nllb.translate(sentence, source=LANGUAGES[self.source_lang], target=LANGUAGES[self.target_lang])
translated_transcript += translated_sentence + splitter + ' '
# print('Text:', sentence)
# print('Text:', translated_sentence)
# print()
translated_transcript = translated_transcript.strip()
return translated_transcript
def get_translated_subtitles(self):
'''
Translates the subtitles into required language
'''
# Creating copy of Transcript Dictionary
subtitles = self.subtitles.copy()
# Creating Instance for Google Translator
gt = GoogleTranslator(source='auto', target=self.target_lang)
for i, s in enumerate(subtitles):
subtitles[i]['text'] = gt.translate(text=s['text'])
return subtitles