Spaces:
Sleeping
Sleeping
import torch | |
import gradio as gr | |
from faster_whisper import WhisperModel | |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline | |
from pydub import AudioSegment | |
import yt_dlp as youtube_dl | |
import tempfile | |
from transformers.pipelines.audio_utils import ffmpeg_read | |
from gradio.components import Audio, Dropdown, Radio, Textbox | |
import os | |
import numpy as np | |
import soundfile as sf | |
os.environ["TOKENIZERS_PARALLELISM"] = "false" | |
# Paramètres | |
FILE_LIMIT_MB = 1000 | |
YT_LENGTH_LIMIT_S = 3600 # Limite de 1 heure pour les vidéos YouTube | |
# Charger les codes de langue | |
from flores200_codes import flores_codes | |
# Fonction pour déterminer le device | |
def set_device(): | |
return torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
device = set_device() | |
# Charger les modèles une seule fois | |
model_dict = {} | |
def load_models(): | |
global model_dict | |
if not model_dict: | |
model_name_dict = { | |
#'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B', | |
'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M', | |
#'nllb-1.3B': 'facebook/nllb-200-1.3B', | |
#'nllb-distilled-1.3B': 'facebook/nllb-200-distilled-1.3B', | |
#'nllb-3.3B': 'facebook/nllb-200-3.3B', | |
# 'nllb-distilled-600M': 'facebook/nllb-200-distilled-600M', | |
} | |
for call_name, real_name in model_name_dict.items(): | |
model = AutoModelForSeq2SeqLM.from_pretrained(real_name) | |
tokenizer = AutoTokenizer.from_pretrained(real_name) | |
model_dict[call_name+'_model'] = model | |
model_dict[call_name+'_tokenizer'] = tokenizer | |
load_models() | |
model_size = "large-v2" | |
model = WhisperModel(model_size) | |
# Fonction pour la transcription | |
def transcribe_audio(audio_file): | |
# model_size = "large-v2" | |
# model = WhisperModel(model_size) | |
# model = WhisperModel(model_size, device=device, compute_type="int8") | |
global model | |
segments, _ = model.transcribe(audio_file, beam_size=1) | |
transcriptions = [("[%.2fs -> %.2fs]" % (seg.start, seg.end), seg.text) for seg in segments] | |
return transcriptions | |
# Fonction pour la traduction | |
def traduction(text, source_lang, target_lang): | |
# Vérifier si les codes de langue sont dans flores_codes | |
if source_lang not in flores_codes or target_lang not in flores_codes: | |
print(f"Code de langue non trouvé : {source_lang} ou {target_lang}") | |
return "" | |
src_code = flores_codes[source_lang] | |
tgt_code = flores_codes[target_lang] | |
model_name = "nllb-distilled-600M" | |
model = model_dict[model_name + "_model"] | |
tokenizer = model_dict[model_name + "_tokenizer"] | |
translator = pipeline("translation", model=model, tokenizer=tokenizer) | |
return translator(text, src_lang=src_code, tgt_lang=tgt_code)[0]["translation_text"] | |
# Fonction principale | |
def full_transcription_and_translation(audio_input, source_lang, target_lang): | |
# Si audio_input est une URL | |
if isinstance(audio_input, str) and audio_input.startswith("http"): | |
audio_file = download_yt_audio(audio_input) | |
# Si audio_input est un dictionnaire contenant des données audio | |
elif isinstance(audio_input, dict) and "array" in audio_input and "sampling_rate" in audio_input: | |
audio_array = audio_input["array"] | |
sampling_rate = audio_input["sampling_rate"] | |
# Écrire le tableau NumPy dans un fichier temporaire WAV | |
with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as f: | |
sf.write(f, audio_array, sampling_rate) | |
audio_file = f.name | |
else: | |
# Supposons que c'est un chemin de fichier | |
audio_file = audio_input | |
transcriptions = transcribe_audio(audio_file) | |
translations = [(timestamp, traduction(text, source_lang, target_lang)) for timestamp, text in transcriptions] | |
# Supprimez le fichier temporaire s'il a été créé | |
if isinstance(audio_input, dict): | |
os.remove(audio_file) | |
return transcriptions, translations | |
# Téléchargement audio YouTube | |
"""def download_yt_audio(yt_url): | |
with tempfile.NamedTemporaryFile(suffix='.mp3') as f: | |
ydl_opts = { | |
'format': 'bestaudio/best', | |
'outtmpl': f.name, | |
'postprocessors': [{ | |
'key': 'FFmpegExtractAudio', | |
'preferredcodec': 'mp3', | |
'preferredquality': '192', | |
}], | |
} | |
with youtube_dl.YoutubeDL(ydl_opts) as ydl: | |
ydl.download([yt_url]) | |
return f.name""" | |
lang_codes = list(flores_codes.keys()) | |
# Interface Gradio | |
def gradio_interface(audio_file, source_lang, target_lang): | |
if audio_file.startswith("http"): | |
audio_file = download_yt_audio(audio_file) | |
transcriptions, translations = full_transcription_and_translation(audio_file, source_lang, target_lang) | |
transcribed_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in transcriptions]) | |
translated_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in translations]) | |
return transcribed_text, translated_text | |
def _return_yt_html_embed(yt_url): | |
video_id = yt_url.split("?v=")[-1] | |
HTML_str = ( | |
f'<center> <iframe width="500" height="320" src="https://www.youtube.com/embed/{video_id}"> </iframe>' | |
" </center>" | |
) | |
return HTML_str | |
def download_yt_audio(yt_url, filename): | |
info_loader = youtube_dl.YoutubeDL() | |
try: | |
info = info_loader.extract_info(yt_url, download=False) | |
except youtube_dl.utils.DownloadError as err: | |
raise gr.Error(str(err)) | |
file_length = info["duration_string"] | |
file_h_m_s = file_length.split(":") | |
file_h_m_s = [int(sub_length) for sub_length in file_h_m_s] | |
if len(file_h_m_s) == 1: | |
file_h_m_s.insert(0, 0) | |
if len(file_h_m_s) == 2: | |
file_h_m_s.insert(0, 0) | |
file_length_s = file_h_m_s[0] * 3600 + file_h_m_s[1] * 60 + file_h_m_s[2] | |
if file_length_s > YT_LENGTH_LIMIT_S: | |
yt_length_limit_hms = time.strftime("%HH:%MM:%SS", time.gmtime(YT_LENGTH_LIMIT_S)) | |
file_length_hms = time.strftime("%HH:%MM:%SS", time.gmtime(file_length_s)) | |
raise gr.Error(f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video.") | |
ydl_opts = {"outtmpl": filename, "format": "worstvideo[ext=mp4]+bestaudio[ext=m4a]/best[ext=mp4]/best"} | |
with youtube_dl.YoutubeDL(ydl_opts) as ydl: | |
try: | |
ydl.download([yt_url]) | |
except youtube_dl.utils.ExtractorError as err: | |
raise gr.Error(str(err)) | |
def yt_transcribe(yt_url, task, max_filesize=75.0): | |
html_embed_str = _return_yt_html_embed(yt_url) | |
global model # S'assurer que le modèle est accessible | |
with tempfile.TemporaryDirectory() as tmpdirname: | |
filepath = os.path.join(tmpdirname, "video.mp4") | |
download_yt_audio(yt_url, filepath) | |
with open(filepath, "rb") as f: | |
inputs = f.read() | |
inputs = ffmpeg_read(inputs, model.feature_extractor.sampling_rate) | |
inputs = {"array": inputs, "sampling_rate": model.feature_extractor.sampling_rate} | |
transcriptions, translations = full_transcription_and_translation(inputs, source_lang, target_lang) | |
transcribed_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in transcriptions]) | |
translated_text = '\n'.join([f"{timestamp}: {text}" for timestamp, text in translations]) | |
return html_embed_str, transcribed_text, translated_text | |
# Interfaces | |
demo = gr.Blocks() | |
with demo: | |
with gr.Tab("Microphone"): | |
gr.Interface( | |
fn=gradio_interface, | |
inputs=[ | |
gr.Audio(sources=["microphone"], type="filepath"), | |
gr.Dropdown(lang_codes, value='French', label='Source Language'), | |
gr.Dropdown(lang_codes, value='English', label='Target Language')], | |
outputs=[gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")] | |
) | |
with gr.Tab("Audio file"): | |
gr.Interface( | |
fn=gradio_interface, | |
inputs=[ | |
gr.Audio(type="filepath", label="Audio file"), | |
gr.Dropdown(lang_codes, value='French', label='Source Language'), | |
gr.Dropdown(lang_codes, value='English', label='Target Language')], | |
outputs=[gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")] | |
) | |
with gr.Tab("YouTube"): | |
gr.Interface( | |
fn=yt_transcribe, | |
inputs=[ | |
gr.Textbox(lines=1, placeholder="Paste the URL to a YouTube video here", label="YouTube URL"), | |
gr.Dropdown(lang_codes, value='French', label='Source Language'), | |
gr.Dropdown(lang_codes, value='English', label='Target Language') | |
], | |
outputs=["html", gr.Textbox(label="Transcribed Text"), gr.Textbox(label="Translated Text")] | |
) | |
#with demo: | |
#gr.TabbedInterface([mf_transcribe, file_transcribe, yt_transcribe], ["Microphone", "Audio file", "YouTube"]) | |
demo.launch() |