app.py

import os
import json
import time
from datetime import datetime
from pathlib import Path
from uuid import uuid4
import tempfile

import gradio as gr
import yt_dlp as youtube_dl
from huggingface_hub import CommitScheduler
from transformers import (
    BitsAndBytesConfig,
    AutoModelForSpeechSeq2Seq,
    AutoTokenizer,
    AutoFeatureExtractor,
    pipeline,
)
from transformers.pipelines.audio_utils import ffmpeg_read

import torch  # If you're using PyTorch
import spaces

os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"

MODEL_NAME = "openai/whisper-large-v3"
BATCH_SIZE = 8
YT_LENGTH_LIMIT_S = 4800  # 1 hour 20 minutes

# Quantization

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

model = AutoModelForSpeechSeq2Seq.from_pretrained(
    MODEL_NAME,
    quantization_config=bnb_config,
    use_cache=False,
    device_map="auto"
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME)

# bnb_config = bnb.QuantizationConfig(bits=4)
pipe = pipeline(
    task="automatic-speech-recognition",
    model=model,
    tokenizer=tokenizer,
    feature_extractor=feature_extractor,
    chunk_length_s=30,
    # device=device,
)

# Define paths and create directory if not exists
JSON_DATASET_DIR = Path("json_dataset")
JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)
JSON_DATASET_PATH = JSON_DATASET_DIR / f"transcriptions-{uuid4()}.json"

# Initialize CommitScheduler for saving data to Hugging Face Dataset
scheduler = CommitScheduler(
    repo_id="transcript-dataset-repo",
    repo_type="dataset",
    folder_path=JSON_DATASET_DIR,
    path_in_repo="data",
)

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"]
    if file_length > YT_LENGTH_LIMIT_S:
        yt_length_limit_hms = time.strftime("%H:%M:%S", time.gmtime(YT_LENGTH_LIMIT_S))
        file_length_hms = time.strftime("%H:%M:%S", time.gmtime(file_length))
        raise gr.Error(
            f"Maximum YouTube length is {yt_length_limit_hms}, got {file_length_hms} YouTube video."
        )
    ydl_opts = {"outtmpl": filename, "format": "bestaudio/best"}
    with youtube_dl.YoutubeDL(ydl_opts) as ydl:
        ydl.download([yt_url])


@spaces.GPU
def yt_transcribe(yt_url, task):
    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, pipe.feature_extractor.sampling_rate)
    inputs = {"array": inputs, "sampling_rate": pipe.feature_extractor.sampling_rate}
    text = pipe(
        inputs,
        batch_size=BATCH_SIZE,
        generate_kwargs={"task": task},
        return_timestamps=True,
    )["text"]
    save_transcription(yt_url, text)
    return text


def save_transcription(yt_url, transcription):
    with scheduler.lock:
        with JSON_DATASET_PATH.open("a") as f:
            json.dump(
                {
                    "url": yt_url,
                    "transcription": transcription,
                    "datetime": datetime.now().isoformat(),
                },
                f,
            )
            f.write("\n")

demo = gr.Blocks()

yt_transcribe_interface = gr.Interface(
    fn=yt_transcribe,
    inputs=[
        gr.Textbox(
            lines=1,
            placeholder="Paste the URL to a YouTube video here",
            label="YouTube URL",
        ),
        gr.Radio(["transcribe", "translate"], label="Task", value="transcribe"),
    ],
    outputs="text",
    title="Whisper Large V3: Transcribe YouTube",
    description=(
        "Transcribe long-form YouTube videos with the click of a button! Demo uses the checkpoint"
        f" [{MODEL_NAME}](https://huggingface.co/{MODEL_NAME}) and 🤗 Transformers to transcribe video files of"
        " arbitrary length."
    ),
    allow_flagging="never",
)

with demo:
    gr.TabbedInterface(
        [yt_transcribe_interface], ["YouTube"]
    )

demo.queue().launch()